HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.32 Apple MPS device Autobatch handling (#13568)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
alexwine36 2024-06-13 18:07:50 -06:00 • committed by GitHub
parent 39fd2661f8
commit e5477713dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/utils/autobatch.py
View file

@ -45,8 +45,8 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
prefix = colorstr("AutoBatch: ")
LOGGER.info(f"{prefix}Computing optimal batch size for imgsz={imgsz} at {fraction * 100}% CUDA memory utilization.")
device = next(model.parameters()).device # get model device
if device.type == "cpu":
LOGGER.info(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}")
if device.type in {"cpu", "mps"}:
LOGGER.info(f"{prefix} ⚠️ intended for CUDA devices, using default batch-size {batch_size}")
return batch_size
if torch.backends.cudnn.benchmark:
LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}")