mirror of
https://github.com/m5stack/yolov5.git
synced 2026-05-20 11:04:37 -07:00
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:
+2
-2
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user