Revert ultralytics-thop to optional (#13290)

This commit is contained in:
Glenn Jocher 2024-06-01 04:20:45 +02:00 committed by GitHub
parent 7453753544
commit 3f2954fbf8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 16 additions and 5 deletions

View file

@ -4,7 +4,6 @@ import contextlib
from copy import deepcopy
from pathlib import Path
import thop
import torch
import torch.nn as nn
@ -66,6 +65,11 @@ 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."""
@ -153,7 +157,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 # GFLOPs
flops = thop.profile(m, inputs=[x.copy() if c else x], verbose=False)[0] / 1e9 * 2 if thop else 0 # GFLOPs
t = time_sync()
for _ in range(10):
m(x.copy() if c else x)