From 3c64d891043643cede117c8e54e30e35aecf2e56 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 29 Aug 2022 20:06:00 +0200 Subject: [PATCH] AutoBatch protect from extreme batch sizes (#9209) If < 1 or > 1024 set output to default batch size 16. May partially address https://github.com/ultralytics/yolov5/issues/9156 Signed-off-by: Glenn Jocher Signed-off-by: Glenn Jocher --- utils/autobatch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/autobatch.py b/utils/autobatch.py index 8d12e46..0115205 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -60,8 +60,8 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16): i = results.index(None) # first fail index if b >= batch_sizes[i]: # y intercept above failure point b = batch_sizes[max(i - 1, 0)] # select prior safe point - if b < 1: # zero or negative batch size - b = 16 + if b < 1 or b > 1024: # b outside of safe range + b = batch_size LOGGER.warning(f'{prefix}WARNING: ⚠️ CUDA anomaly detected, recommend restart environment and retry command.') fraction = np.polyval(p, b) / t # actual fraction predicted