AutoBatch checks against failed solutions (#8159)

utils/autobatch.py

@@ -8,7 +8,7 @@ from copy import deepcopy
 import numpy as np
 import torch
 
-from utils.general import LOGGER, colorstr
+from utils.general import LOGGER, colorstr, emojis
 from utils.torch_utils import profile
 
 
@@ -26,6 +26,7 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
     #     model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
     #     print(autobatch(model))
 
+    # Check device
     prefix = colorstr('AutoBatch: ')
     LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}')
     device = next(model.parameters()).device  # get model device
@@ -33,25 +34,33 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
         LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}')
         return batch_size
 
+    # Inspect CUDA memory
     gb = 1 << 30  # bytes to GiB (1024 ** 3)
     d = str(device).upper()  # 'CUDA:0'
     properties = torch.cuda.get_device_properties(device)  # device properties
-    t = properties.total_memory / gb  # (GiB)
-    r = torch.cuda.memory_reserved(device) / gb  # (GiB)
-    a = torch.cuda.memory_allocated(device) / gb  # (GiB)
-    f = t - (r + a)  # free inside reserved
+    t = properties.total_memory / gb  # GiB total
+    r = torch.cuda.memory_reserved(device) / gb  # GiB reserved
+    a = torch.cuda.memory_allocated(device) / gb  # GiB allocated
+    f = t - (r + a)  # GiB free
     LOGGER.info(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free')
 
+    # Profile batch sizes
     batch_sizes = [1, 2, 4, 8, 16]
     try:
         img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes]
-        y = profile(img, model, n=3, device=device)
+        results = profile(img, model, n=3, device=device)
     except Exception as e:
         LOGGER.warning(f'{prefix}{e}')
 
-    y = [x[2] for x in y if x]  # memory [2]
-    batch_sizes = batch_sizes[:len(y)]
-    p = np.polyfit(batch_sizes, y, deg=1)  # first degree polynomial fit
+    # Fit a solution
+    y = [x[2] for x in results if x]  # memory [2]
+    p = np.polyfit(batch_sizes[:len(y)], y, deg=1)  # first degree polynomial fit
     b = int((f * fraction - p[1]) / p[0])  # y intercept (optimal batch size)
-    LOGGER.info(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)')
+    if None in results:  # some sizes failed
+        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
+
+    fraction = np.polyval(p, b) / t  # actual fraction predicted
+    LOGGER.info(emojis(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅'))
     return b