zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

FLOPS computation device bug fix (#1447) 

* Update torch_utils.py

fix issue#113 , inputs device should be same with model parameters' device

* Update torch_utils.py

* Update torch_utils.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
yujun 2020-11-19 19:56:20 +08:00 committed by GitHub
parent af8aee76b3
commit 05a955a3f6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
utils/torch_utils.py
View File

@ -153,7 +153,8 @@ def model_info(model, verbose=False, img_size=640):
try: # FLOPS try: # FLOPS
from thop import profile from thop import profile
stride = int(model.stride.max()) stride = int(model.stride.max())
flops = profile(deepcopy(model), inputs=(torch.zeros(1, 3, stride, stride),), verbose=False)[0] / 1E9 * 2 img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device) # input
flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride FLOPS
img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 FLOPS fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 FLOPS
except (ImportError, Exception): except (ImportError, Exception):