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

ultralytics 8.3.31 add max_num_obj factor for AutoBatch (#17514)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Laughing 2024-11-14 06:51:24 +08:00 • committed by GitHub
parent e100484422
commit 4453ddab93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 38 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

10
ultralytics/utils/torch_utils.py
View file

@ -623,7 +623,7 @@ def convert_optimizer_state_dict_to_fp16(state_dict):
return state_dict
def profile(input, ops, n=10, device=None):
def profile(input, ops, n=10, device=None, max_num_obj=0):
"""
Ultralytics speed, memory and FLOPs profiler.
@ -671,6 +671,14 @@ def profile(input, ops, n=10, device=None):
t[2] = float("nan")
tf += (t[1] - t[0]) * 1000 / n # ms per op forward
tb += (t[2] - t[1]) * 1000 / n # ms per op backward
if max_num_obj: # simulate training with predictions per image grid (for AutoBatch)
torch.randn(
x.shape[0],
max_num_obj,
int(sum([(x.shape[-1] / s) * (x.shape[-2] / s) for s in m.stride.tolist()])),
device=device,
dtype=torch.float32,
)
mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0 # (GB)
s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y)) # shapes
p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0 # parameters