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

Use new ultralytics-thop package (#13282)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-06-01 01:51:39 +02:00 • committed by GitHub
parent 8fb140688a
commit 7453753544
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 6 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

8
ultralytics/nn/tasks.py
View file

@ -4,6 +4,7 @@ import contextlib
from copy import deepcopy
from pathlib import Path
import thop
import torch
import torch.nn as nn
@ -65,11 +66,6 @@ from ultralytics.utils.torch_utils import (
time_sync,
)
try:
import thop
except ImportError:
thop = None
class BaseModel(nn.Module):
"""The BaseModel class serves as a base class for all the models in the Ultralytics YOLO family."""
@ -157,7 +153,7 @@ class BaseModel(nn.Module):
None
"""
c = m == self.model[-1] and isinstance(x, list) # is final layer list, copy input as inplace fix
flops = thop.profile(m, inputs=[x.copy() if c else x], verbose=False)[0] / 1e9 * 2 if thop else 0 # FLOPs
flops = thop.profile(m, inputs=[x.copy() if c else x], verbose=False)[0] / 1e9 * 2 # GFLOPs
t = time_sync()
for _ in range(10):
m(x.copy() if c else x)