Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

AutoBatch cudnn.benchmark=True fix (#9448)

Browse Source 
* AutoBatch `cudnn.benchmark=True` fix

May resolve https://github.com/ultralytics/yolov5/issues/9287

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

* Update autobatch.py

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

* Update autobatch.py

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

* Update general.py

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 authored and GitHub committed 2022-09-16 21:46:07 +02:00
1 parent db06f495db
commit 5e1a9553fb
2 files changed
+4 -1
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

utils/autobatch.py
+3
View File
@@ -33,6 +33,9 @@ def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
if device.type == 'cpu':
LOGGER.info(f'{prefix}CUDA not detected, using default CPU 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}')
return batch_size
# Inspect CUDA memory
gb = 1 << 30 # bytes to GiB (1024 ** 3)