From 8154a27e9b332dd2dac7a107003462f0e20f6179 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 8 Oct 2024 21:50:28 +0200 Subject: [PATCH] AutoBatch larger batches on larger GPUs (#16785) --- ultralytics/utils/autobatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/utils/autobatch.py b/ultralytics/utils/autobatch.py index 581fe09027..6a0d9cbc29 100644 --- a/ultralytics/utils/autobatch.py +++ b/ultralytics/utils/autobatch.py @@ -67,7 +67,7 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch): 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] + batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64] try: img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes] results = profile(img, model, n=1, device=device)