From 862cc10a7d854d3a80bbeabf68e79eabb4af47eb Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 19 Sep 2022 12:18:07 +0200 Subject: [PATCH] AutoBatch report include reserved+allocated May resolve https://github.com/ultralytics/yolov5/issues/9287#issuecomment-1250767031 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 49435f51a244..bdeb91c3d2bd 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -19,7 +19,7 @@ def check_train_batch_size(model, imgsz=640, amp=True): def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): - # Automatically estimate best batch size to use `fraction` of available CUDA memory + # Automatically estimate best YOLOv5 batch size to use `fraction` of available CUDA memory # Usage: # import torch # from utils.autobatch import autobatch @@ -67,6 +67,6 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16): 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 + fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅') return b