ultralytics-thop 2.0.14__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- thop/__init__.py +11 -0
- thop/fx_profile.py +237 -0
- thop/profile.py +228 -0
- thop/rnn_hooks.py +202 -0
- thop/utils.py +50 -0
- thop/vision/__init__.py +1 -0
- thop/vision/basic_hooks.py +146 -0
- thop/vision/calc_func.py +127 -0
- ultralytics_thop-2.0.14.dist-info/LICENSE +661 -0
- ultralytics_thop-2.0.14.dist-info/METADATA +199 -0
- ultralytics_thop-2.0.14.dist-info/RECORD +13 -0
- ultralytics_thop-2.0.14.dist-info/WHEEL +5 -0
- ultralytics_thop-2.0.14.dist-info/top_level.txt +1 -0
thop/utils.py
ADDED
@@ -0,0 +1,50 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
from collections.abc import Iterable
|
4
|
+
|
5
|
+
COLOR_RED = "91m"
|
6
|
+
COLOR_GREEN = "92m"
|
7
|
+
COLOR_YELLOW = "93m"
|
8
|
+
|
9
|
+
|
10
|
+
def colorful_print(fn_print, color=COLOR_RED):
|
11
|
+
"""A decorator to print text in the specified terminal color by wrapping the given print function."""
|
12
|
+
|
13
|
+
def actual_call(*args, **kwargs):
|
14
|
+
print(f"\033[{color}", end="")
|
15
|
+
fn_print(*args, **kwargs)
|
16
|
+
print("\033[00m", end="")
|
17
|
+
|
18
|
+
return actual_call
|
19
|
+
|
20
|
+
|
21
|
+
prRed = colorful_print(print, color=COLOR_RED)
|
22
|
+
prGreen = colorful_print(print, color=COLOR_GREEN)
|
23
|
+
prYellow = colorful_print(print, color=COLOR_YELLOW)
|
24
|
+
|
25
|
+
|
26
|
+
def clever_format(nums, format="%.2f"):
|
27
|
+
"""Formats numbers into human-readable strings with units (K for thousand, M for million, etc.)."""
|
28
|
+
if not isinstance(nums, Iterable):
|
29
|
+
nums = [nums]
|
30
|
+
clever_nums = []
|
31
|
+
|
32
|
+
for num in nums:
|
33
|
+
if num > 1e12:
|
34
|
+
clever_nums.append(format % (num / 1e12) + "T")
|
35
|
+
elif num > 1e9:
|
36
|
+
clever_nums.append(format % (num / 1e9) + "G")
|
37
|
+
elif num > 1e6:
|
38
|
+
clever_nums.append(format % (num / 1e6) + "M")
|
39
|
+
elif num > 1e3:
|
40
|
+
clever_nums.append(format % (num / 1e3) + "K")
|
41
|
+
else:
|
42
|
+
clever_nums.append(format % num + "B")
|
43
|
+
|
44
|
+
return clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
|
45
|
+
|
46
|
+
|
47
|
+
if __name__ == "__main__":
|
48
|
+
prRed("hello", "world")
|
49
|
+
prGreen("hello", "world")
|
50
|
+
prYellow("hello", "world")
|
thop/vision/__init__.py
ADDED
@@ -0,0 +1 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
@@ -0,0 +1,146 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
import logging
|
4
|
+
|
5
|
+
import torch.nn as nn
|
6
|
+
from torch.nn.modules.conv import _ConvNd
|
7
|
+
|
8
|
+
from .calc_func import *
|
9
|
+
|
10
|
+
multiply_adds = 1
|
11
|
+
|
12
|
+
|
13
|
+
def count_parameters(m, x, y):
|
14
|
+
"""Calculate and return the total number of learnable parameters in a given PyTorch model."""
|
15
|
+
m.total_params[0] = calculate_parameters(m.parameters())
|
16
|
+
|
17
|
+
|
18
|
+
def zero_ops(m, x, y):
|
19
|
+
"""Incrementally add zero operations to the model's total operations count."""
|
20
|
+
m.total_ops += calculate_zero_ops()
|
21
|
+
|
22
|
+
|
23
|
+
def count_convNd(m: _ConvNd, x, y: torch.Tensor):
|
24
|
+
"""Calculate and add the number of convolutional operations (FLOPs) for a ConvNd layer to the model's total ops."""
|
25
|
+
x = x[0]
|
26
|
+
|
27
|
+
m.total_ops += calculate_conv2d_flops(
|
28
|
+
input_size=list(x.shape),
|
29
|
+
output_size=list(y.shape),
|
30
|
+
kernel_size=list(m.weight.shape),
|
31
|
+
groups=m.groups,
|
32
|
+
bias=m.bias,
|
33
|
+
)
|
34
|
+
# N x Cout x H x W x (Cin x Kw x Kh + bias)
|
35
|
+
# m.total_ops += calculate_conv(
|
36
|
+
# bias_ops,
|
37
|
+
# torch.zeros(m.weight.size()[2:]).numel(),
|
38
|
+
# y.nelement(),
|
39
|
+
# m.in_channels,
|
40
|
+
# m.groups,
|
41
|
+
# )
|
42
|
+
|
43
|
+
|
44
|
+
def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor):
|
45
|
+
"""Calculates and updates total operations (FLOPs) for a convolutional layer in a PyTorch model."""
|
46
|
+
x = x[0]
|
47
|
+
|
48
|
+
# N x H x W (exclude Cout)
|
49
|
+
output_size = torch.zeros(y.size()[:1] + y.size()[2:]).numel()
|
50
|
+
# # Cout x Cin x Kw x Kh
|
51
|
+
# kernel_ops = m.weight.nelement()
|
52
|
+
# if m.bias is not None:
|
53
|
+
# # Cout x 1
|
54
|
+
# kernel_ops += + m.bias.nelement()
|
55
|
+
# # x N x H x W x Cout x (Cin x Kw x Kh + bias)
|
56
|
+
# m.total_ops += torch.DoubleTensor([int(output_size * kernel_ops)])
|
57
|
+
m.total_ops += calculate_conv(m.bias.nelement(), m.weight.nelement(), output_size)
|
58
|
+
|
59
|
+
|
60
|
+
def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y):
|
61
|
+
"""Calculate and add the FLOPs for a batch normalization layer, including elementwise and affine operations."""
|
62
|
+
# https://github.com/Lyken17/pytorch-OpCounter/issues/124
|
63
|
+
# y = (x - mean) / sqrt(eps + var) * weight + bias
|
64
|
+
x = x[0]
|
65
|
+
# bn is by default fused in inference
|
66
|
+
flops = calculate_norm(x.numel())
|
67
|
+
if getattr(m, "affine", False) or getattr(m, "elementwise_affine", False):
|
68
|
+
flops *= 2
|
69
|
+
m.total_ops += flops
|
70
|
+
|
71
|
+
|
72
|
+
# def count_layer_norm(m, x, y):
|
73
|
+
# x = x[0]
|
74
|
+
# m.total_ops += calculate_norm(x.numel())
|
75
|
+
|
76
|
+
|
77
|
+
# def count_instance_norm(m, x, y):
|
78
|
+
# x = x[0]
|
79
|
+
# m.total_ops += calculate_norm(x.numel())
|
80
|
+
|
81
|
+
|
82
|
+
def count_prelu(m, x, y):
|
83
|
+
"""Calculate and update the total operation counts for a PReLU layer using input element number."""
|
84
|
+
x = x[0]
|
85
|
+
|
86
|
+
nelements = x.numel()
|
87
|
+
if not m.training:
|
88
|
+
m.total_ops += calculate_relu(nelements)
|
89
|
+
|
90
|
+
|
91
|
+
def count_relu(m, x, y):
|
92
|
+
"""Calculate and update the total operation counts for a ReLU layer."""
|
93
|
+
x = x[0]
|
94
|
+
m.total_ops += calculate_relu_flops(list(x.shape))
|
95
|
+
|
96
|
+
|
97
|
+
def count_softmax(m, x, y):
|
98
|
+
"""Calculate and update the total operation counts for a Softmax layer in a PyTorch model."""
|
99
|
+
x = x[0]
|
100
|
+
nfeatures = x.size()[m.dim]
|
101
|
+
batch_size = x.numel() // nfeatures
|
102
|
+
|
103
|
+
m.total_ops += calculate_softmax(batch_size, nfeatures)
|
104
|
+
|
105
|
+
|
106
|
+
def count_avgpool(m, x, y):
|
107
|
+
"""Calculate and update the total number of operations (FLOPs) for an AvgPool layer based on the output elements."""
|
108
|
+
# total_div = 1
|
109
|
+
# kernel_ops = total_add + total_div
|
110
|
+
num_elements = y.numel()
|
111
|
+
m.total_ops += calculate_avgpool(num_elements)
|
112
|
+
|
113
|
+
|
114
|
+
def count_adap_avgpool(m, x, y):
|
115
|
+
"""Calculate and update the total operation counts for an AdaptiveAvgPool layer using kernel and element counts."""
|
116
|
+
kernel = torch.div(torch.DoubleTensor([*(x[0].shape[2:])]), torch.DoubleTensor([*(y.shape[2:])]))
|
117
|
+
total_add = torch.prod(kernel)
|
118
|
+
num_elements = y.numel()
|
119
|
+
m.total_ops += calculate_adaptive_avg(total_add, num_elements)
|
120
|
+
|
121
|
+
|
122
|
+
# TODO: verify the accuracy
|
123
|
+
def count_upsample(m, x, y):
|
124
|
+
"""Update total operations counter for upsampling layers based on the mode used."""
|
125
|
+
if m.mode not in (
|
126
|
+
"nearest",
|
127
|
+
"linear",
|
128
|
+
"bilinear",
|
129
|
+
"bicubic",
|
130
|
+
): # "trilinear"
|
131
|
+
logging.warning(f"mode {m.mode} is not implemented yet, take it a zero op")
|
132
|
+
m.total_ops += 0
|
133
|
+
else:
|
134
|
+
x = x[0]
|
135
|
+
m.total_ops += calculate_upsample(m.mode, y.nelement())
|
136
|
+
|
137
|
+
|
138
|
+
# nn.Linear
|
139
|
+
def count_linear(m, x, y):
|
140
|
+
"""Counts total operations for nn.Linear layers using input and output element dimensions."""
|
141
|
+
total_mul = m.in_features
|
142
|
+
# total_add = m.in_features - 1
|
143
|
+
# total_add += 1 if m.bias is not None else 0
|
144
|
+
num_elements = y.numel()
|
145
|
+
|
146
|
+
m.total_ops += calculate_linear(total_mul, num_elements)
|
thop/vision/calc_func.py
ADDED
@@ -0,0 +1,127 @@
|
|
1
|
+
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
|
2
|
+
|
3
|
+
import warnings
|
4
|
+
|
5
|
+
import numpy as np
|
6
|
+
import torch
|
7
|
+
|
8
|
+
|
9
|
+
def l_prod(in_list):
|
10
|
+
"""Compute the product of all elements in the input list."""
|
11
|
+
res = 1
|
12
|
+
for _ in in_list:
|
13
|
+
res *= _
|
14
|
+
return res
|
15
|
+
|
16
|
+
|
17
|
+
def l_sum(in_list):
|
18
|
+
"""Calculate the sum of all numerical elements in a list."""
|
19
|
+
return sum(in_list)
|
20
|
+
|
21
|
+
|
22
|
+
def calculate_parameters(param_list):
|
23
|
+
"""Calculate the total number of parameters in a list of tensors using the product of their shapes."""
|
24
|
+
return sum(torch.DoubleTensor([p.nelement()]) for p in param_list)
|
25
|
+
|
26
|
+
|
27
|
+
def calculate_zero_ops():
|
28
|
+
"""Initializes and returns a tensor with all elements set to zero."""
|
29
|
+
return torch.DoubleTensor([0])
|
30
|
+
|
31
|
+
|
32
|
+
def calculate_conv2d_flops(input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False):
|
33
|
+
"""Calculate FLOPs for a Conv2D layer using input/output sizes, kernel size, groups, and the bias flag."""
|
34
|
+
# n, in_c, ih, iw = input_size
|
35
|
+
# out_c, in_c, kh, kw = kernel_size
|
36
|
+
in_c = input_size[1]
|
37
|
+
g = groups
|
38
|
+
return l_prod(output_size) * (in_c // g) * l_prod(kernel_size[2:])
|
39
|
+
|
40
|
+
|
41
|
+
def calculate_conv(bias, kernel_size, output_size, in_channel, group):
|
42
|
+
"""Calculate FLOPs for convolutional layers given bias, kernel size, output size, in_channels, and groups."""
|
43
|
+
warnings.warn("This API is being deprecated.")
|
44
|
+
return torch.DoubleTensor([output_size * (in_channel / group * kernel_size + bias)])
|
45
|
+
|
46
|
+
|
47
|
+
def calculate_norm(input_size):
|
48
|
+
"""Compute the L2 norm of a tensor or array based on its input size."""
|
49
|
+
return torch.DoubleTensor([2 * input_size])
|
50
|
+
|
51
|
+
|
52
|
+
def calculate_relu_flops(input_size):
|
53
|
+
"""Calculates the FLOPs for a ReLU activation function based on the input tensor's dimensions."""
|
54
|
+
return 0
|
55
|
+
|
56
|
+
|
57
|
+
def calculate_relu(input_size: torch.Tensor):
|
58
|
+
"""Convert an input tensor to a DoubleTensor with the same value (deprecated)."""
|
59
|
+
warnings.warn("This API is being deprecated")
|
60
|
+
return torch.DoubleTensor([int(input_size)])
|
61
|
+
|
62
|
+
|
63
|
+
def calculate_softmax(batch_size, nfeatures):
|
64
|
+
"""Compute FLOPs for a softmax activation given batch size and feature count."""
|
65
|
+
total_exp = nfeatures
|
66
|
+
total_add = nfeatures - 1
|
67
|
+
total_div = nfeatures
|
68
|
+
total_ops = batch_size * (total_exp + total_add + total_div)
|
69
|
+
return torch.DoubleTensor([int(total_ops)])
|
70
|
+
|
71
|
+
|
72
|
+
def calculate_avgpool(input_size):
|
73
|
+
"""Calculate the average pooling size for a given input tensor."""
|
74
|
+
return torch.DoubleTensor([int(input_size)])
|
75
|
+
|
76
|
+
|
77
|
+
def calculate_adaptive_avg(kernel_size, output_size):
|
78
|
+
"""Calculate FLOPs for adaptive average pooling given kernel size and output size."""
|
79
|
+
total_div = 1
|
80
|
+
kernel_op = kernel_size + total_div
|
81
|
+
return torch.DoubleTensor([int(kernel_op * output_size)])
|
82
|
+
|
83
|
+
|
84
|
+
def calculate_upsample(mode: str, output_size):
|
85
|
+
"""Calculate the operations required for various upsample methods based on mode and output size."""
|
86
|
+
total_ops = output_size
|
87
|
+
if mode == "bicubic":
|
88
|
+
total_ops *= 224 + 35
|
89
|
+
elif mode == "bilinear":
|
90
|
+
total_ops *= 11
|
91
|
+
elif mode == "linear":
|
92
|
+
total_ops *= 5
|
93
|
+
elif mode == "trilinear":
|
94
|
+
total_ops *= 13 * 2 + 5
|
95
|
+
return torch.DoubleTensor([int(total_ops)])
|
96
|
+
|
97
|
+
|
98
|
+
def calculate_linear(in_feature, num_elements):
|
99
|
+
"""Calculate the linear operation count for given input feature and number of elements."""
|
100
|
+
return torch.DoubleTensor([int(in_feature * num_elements)])
|
101
|
+
|
102
|
+
|
103
|
+
def counter_matmul(input_size, output_size):
|
104
|
+
"""Calculate the total number of operations for matrix multiplication given input and output sizes."""
|
105
|
+
input_size = np.array(input_size)
|
106
|
+
output_size = np.array(output_size)
|
107
|
+
return np.prod(input_size) * output_size[-1]
|
108
|
+
|
109
|
+
|
110
|
+
def counter_mul(input_size):
|
111
|
+
"""Calculate the total number of operations for element-wise multiplication given the input size."""
|
112
|
+
return input_size
|
113
|
+
|
114
|
+
|
115
|
+
def counter_pow(input_size):
|
116
|
+
"""Computes the total scalar multiplications required for power operations based on input size."""
|
117
|
+
return input_size
|
118
|
+
|
119
|
+
|
120
|
+
def counter_sqrt(input_size):
|
121
|
+
"""Calculate the total number of scalar operations required for a square root operation given an input size."""
|
122
|
+
return input_size
|
123
|
+
|
124
|
+
|
125
|
+
def counter_div(input_size):
|
126
|
+
"""Calculate the total number of scalar operations for a division operation given an input size."""
|
127
|
+
return input_size
|