From 74537535444edbfaa1479754632995bc4dbf1097 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 1 Jun 2024 01:51:39 +0200
Subject: [PATCH] Use new `ultralytics-thop` package (#13282)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
pyproject.toml | 4 ++--
ultralytics/nn/tasks.py | 8 ++------
ultralytics/utils/torch_utils.py | 11 ++---------
3 files changed, 6 insertions(+), 17 deletions(-)
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