diff --git a/pyproject.toml b/pyproject.toml index 49296f77f0..e21c4b3e41 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,9 +75,9 @@ dependencies = [ "tqdm>=4.64.0", # progress bars "psutil", # system utilization "py-cpuinfo", # display CPU info - "thop>=0.1.1", # FLOPs computation "pandas>=1.1.4", "seaborn>=0.11.0", # plotting + "ultralytics-thop>=0.2.4", # FLOPs computation https://github.com/ultralytics/thop ] # Optional dependencies ------------------------------------------------------------------------------------------------ @@ -94,7 +94,7 @@ dev = [ "mkdocstrings[python]", "mkdocs-jupyter", # for notebooks "mkdocs-redirects", # for 301 redirects - "mkdocs-ultralytics-plugin>=0.0.44", # for meta descriptions and images, dates and authors + "mkdocs-ultralytics-plugin>=0.0.45", # for meta descriptions and images, dates and authors ] export = [ "onnx>=1.12.0", # ONNX export diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 52e999c92a..3e46ded7eb 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -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) diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py index 919fee07a2..751e98ba4c 100644 --- a/ultralytics/utils/torch_utils.py +++ b/ultralytics/utils/torch_utils.py @@ -11,6 +11,7 @@ from pathlib import Path from typing import Union import numpy as np +import thop import torch import torch.distributed as dist import torch.nn as nn @@ -27,11 +28,6 @@ from ultralytics.utils import ( ) from ultralytics.utils.checks import check_version -try: - import thop -except ImportError: - thop = None - # Version checks (all default to version>=min_version) TORCH_1_9 = check_version(torch.__version__, "1.9.0") TORCH_1_13 = check_version(torch.__version__, "1.13.0") @@ -308,9 +304,6 @@ def model_info_for_loggers(trainer): def get_flops(model, imgsz=640): """Return a YOLO model's FLOPs.""" - if not thop: - return 0.0 # if not installed return 0.0 GFLOPs - try: model = de_parallel(model) p = next(model.parameters()) @@ -571,7 +564,7 @@ def profile(input, ops, n=10, device=None): m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m tf, tb, t = 0, 0, [0, 0, 0] # dt forward, backward try: - flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1e9 * 2 if thop else 0 # GFLOPs + flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1e9 * 2 # GFLOPs except Exception: flops = 0