AutoBatch report include reserved+allocated (#9491)

May resolve https://github.com/ultralytics/yolov5/issues/9287#issuecomment-1250767031

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher
2022-09-19 12:18:55 +02:00
committed by GitHub
parent fda8aa551d
commit f038ad7172
+2 -2
View File
@@ -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