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

AutoBatch improve cache clearing (#16744)

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Ultralytics Assistant 2024-10-07 20:11:10 +02:00 • committed by GitHub
parent f7808dc61b
commit 97521186a7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 8 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
ultralytics/utils/autobatch.py
View file

@ -69,7 +69,7 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
batch_sizes = [1, 2, 4, 8, 16]
try:
img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
results = profile(img, model, n=3, device=device)
results = profile(img, model, n=1, device=device)
# Fit a solution
y = [x[2] for x in results if x] # memory [2]
@ -89,3 +89,5 @@ def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
except Exception as e:
LOGGER.warning(f"{prefix}WARNING ⚠️ error detected: {e}, using default batch-size {batch_size}.")
return batch_size
finally:
torch.cuda.empty_cache()

4
ultralytics/utils/torch_utils.py
View file

@ -643,7 +643,8 @@ def profile(input, ops, n=10, device=None):
f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"
f"{'input':>24s}{'output':>24s}"
)
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache()
for x in input if isinstance(input, list) else [input]:
x = x.to(device)
x.requires_grad = True
@ -677,6 +678,7 @@ def profile(input, ops, n=10, device=None):
except Exception as e:
LOGGER.info(e)
results.append(None)
finally:
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache()
return results