From b50a327a043057103e995ce7d472a4804f295a0f Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 19 Feb 2025 23:01:56 +0800 Subject: [PATCH] Allow missing `thop` package (#19314) Co-authored-by: UltralyticsAssistant --- ultralytics/nn/tasks.py | 6 +++++- ultralytics/utils/torch_utils.py | 11 +++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/ultralytics/nn/tasks.py b/ultralytics/nn/tasks.py index 4ed8d7da39..4d845ae9c5 100644 --- a/ultralytics/nn/tasks.py +++ b/ultralytics/nn/tasks.py @@ -7,7 +7,6 @@ import types from copy import deepcopy from pathlib import Path -import thop import torch from ultralytics.nn.modules import ( @@ -86,6 +85,11 @@ from ultralytics.utils.torch_utils import ( time_sync, ) +try: + import thop +except ImportError: + thop = None # conda support without 'ultralytics-thop' installed + class BaseModel(torch.nn.Module): """The BaseModel class serves as a base class for all the models in the Ultralytics YOLO family.""" diff --git a/ultralytics/utils/torch_utils.py b/ultralytics/utils/torch_utils.py index a589b7d640..83eb303725 100644 --- a/ultralytics/utils/torch_utils.py +++ b/ultralytics/utils/torch_utils.py @@ -12,7 +12,6 @@ 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 @@ -31,6 +30,11 @@ from ultralytics.utils import ( ) from ultralytics.utils.checks import check_version +try: + import thop +except ImportError: + thop = None # conda support without 'ultralytics-thop' installed + # 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") @@ -370,6 +374,9 @@ 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()) @@ -681,7 +688,7 @@ def profile(input, ops, n=10, device=None, max_num_obj=0): 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(deepcopy(m), inputs=[x], verbose=False)[0] / 1e9 * 2 # GFLOPs + flops = thop.profile(deepcopy(m), inputs=[x], verbose=False)[0] / 1e9 * 2 if thop else 0 # GFLOPs except Exception: flops = 0