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

8
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,8 +678,9 @@ def profile(input, ops, n=10, device=None):
except Exception as e:
LOGGER.info(e)
results.append(None)
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache()
finally:
gc.collect() # attempt to free unused memory
torch.cuda.empty_cache()
return results