Unverified 提交 f038ad71 authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

AutoBatch report include reserved+allocated (#9491)

上级 fda8aa55
......@@ -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
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论