AutoBatch report include reserved+allocated

May resolve #9287 (comment) Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
ultralytics · Sep 19, 2022 · 862cc10 · 862cc10
1 parent fda8aa5
commit 862cc10
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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