Use new ultralytics-thop package (#13282)
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
parent
8fb140688a
commit
7453753544
3 changed files with 6 additions and 17 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue