ultralytics-thop 2.0.15__tar.gz → 2.0.17__tar.gz

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.
Files changed (23) hide show
  1. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/PKG-INFO +4 -4
  2. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/README.md +3 -3
  3. ultralytics_thop-2.0.17/tests/test_conv2d.py +219 -0
  4. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/tests/test_matmul.py +1 -1
  5. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/__init__.py +3 -2
  6. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/fx_profile.py +9 -8
  7. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/profile.py +28 -5
  8. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/vision/basic_hooks.py +30 -1
  9. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/vision/calc_func.py +9 -4
  10. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/ultralytics_thop.egg-info/PKG-INFO +4 -4
  11. ultralytics_thop-2.0.15/tests/test_conv2d.py +0 -61
  12. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/LICENSE +0 -0
  13. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/pyproject.toml +0 -0
  14. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/setup.cfg +0 -0
  15. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/tests/test_relu.py +0 -0
  16. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/tests/test_utils.py +0 -0
  17. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/rnn_hooks.py +0 -0
  18. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/utils.py +0 -0
  19. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/thop/vision/__init__.py +0 -0
  20. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/ultralytics_thop.egg-info/SOURCES.txt +0 -0
  21. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/ultralytics_thop.egg-info/dependency_links.txt +0 -0
  22. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/ultralytics_thop.egg-info/requires.txt +0 -0
  23. {ultralytics_thop-2.0.15 → ultralytics_thop-2.0.17}/ultralytics_thop.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ultralytics-thop
3
- Version: 2.0.15
3
+ Version: 2.0.17
4
4
  Summary: Ultralytics THOP package for fast computation of PyTorch model FLOPs and parameters.
5
5
  Author-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>
6
6
  Maintainer-email: Ultralytics <hello@ultralytics.com>
@@ -56,7 +56,7 @@ THOP offers an intuitive API designed to profile PyTorch models by calculating t
56
56
 
57
57
  Get started with THOP quickly by installing it via pip:
58
58
 
59
- [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://www.pepy.tech/projects/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
59
+ [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://clickpy.clickhouse.com/dashboard/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
60
60
 
61
61
  ```bash
62
62
  pip install ultralytics-thop
@@ -74,7 +74,7 @@ This ensures you have the most recent version, incorporating the latest improvem
74
74
 
75
75
  ### Basic Usage
76
76
 
77
- Profiling a standard PyTorch model like [ResNet50](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
77
+ Profiling a standard PyTorch model like [ResNet50](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
78
78
 
79
79
  ```python
80
80
  import torch
@@ -230,7 +230,7 @@ The table below showcases the parameters and MACs for several popular [computer
230
230
 
231
231
  ## 🙌 Contribute
232
232
 
233
- We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
233
+ We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://docs.pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
234
234
 
235
235
  ## 📜 License
236
236
 
@@ -17,7 +17,7 @@ THOP offers an intuitive API designed to profile PyTorch models by calculating t
17
17
 
18
18
  Get started with THOP quickly by installing it via pip:
19
19
 
20
- [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://www.pepy.tech/projects/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
20
+ [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://clickpy.clickhouse.com/dashboard/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
21
21
 
22
22
  ```bash
23
23
  pip install ultralytics-thop
@@ -35,7 +35,7 @@ This ensures you have the most recent version, incorporating the latest improvem
35
35
 
36
36
  ### Basic Usage
37
37
 
38
- Profiling a standard PyTorch model like [ResNet50](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
38
+ Profiling a standard PyTorch model like [ResNet50](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
39
39
 
40
40
  ```python
41
41
  import torch
@@ -191,7 +191,7 @@ The table below showcases the parameters and MACs for several popular [computer
191
191
 
192
192
  ## 🙌 Contribute
193
193
 
194
- We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
194
+ We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://docs.pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
195
195
 
196
196
  ## 📜 License
197
197
 
@@ -0,0 +1,219 @@
1
+ # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+
6
+ from thop import profile
7
+
8
+
9
+ class TestUtils:
10
+ """Utility class for testing Conv2D layers with and without bias, profiling MACs and parameters using THOP."""
11
+
12
+ def test_conv2d_no_bias(self):
13
+ """Tests a 2D Conv layer without bias using THOP profiling on predefined input dimensions and parameters."""
14
+ n, in_c, ih, iw = 1, 3, 32, 32 # torch.randint(1, 10, (4,)).tolist()
15
+ out_c, kh, kw = 12, 5, 5
16
+ s, p, d, g = 1, 1, 1, 1
17
+
18
+ net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False)
19
+ data = torch.randn(n, in_c, ih, iw)
20
+ out = net(data)
21
+
22
+ _, _, oh, ow = out.shape
23
+ print(f"Conv2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
24
+ macs, params = profile(net, inputs=(data,))
25
+ assert macs == 810000, f"{macs} v.s. 810000"
26
+
27
+ def test_conv2d(self):
28
+ """Tests Conv2D layer with bias, profiling MACs and params for specific input dimensions and layer configs."""
29
+ n, in_c, ih, iw = 1, 3, 32, 32 # torch.randint(1, 10, (4,)).tolist()
30
+ out_c, kh, kw = 12, 5, 5
31
+ s, p, d, g = 1, 1, 1, 1
32
+
33
+ net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=True)
34
+ data = torch.randn(n, in_c, ih, iw)
35
+ out = net(data)
36
+
37
+ _, _, oh, ow = out.shape
38
+ print(f"Conv2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
39
+ macs, params = profile(net, inputs=(data,))
40
+ assert macs == 810000, f"{macs} v.s. 810000"
41
+
42
+ def test_conv2d_random(self):
43
+ """Test Conv2D layer with random parameters and validate the computed MACs and parameters using 'profile'."""
44
+ for _ in range(10):
45
+ out_c, kh, kw = torch.randint(1, 20, (3,)).tolist()
46
+ n, in_c, ih, iw = torch.randint(1, 20, (4,)).tolist() # torch.randint(1, 10, (4,)).tolist()
47
+ ih += kh
48
+ iw += kw
49
+ s, p, d, g = 1, 1, 1, 1
50
+
51
+ net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False)
52
+ data = torch.randn(n, in_c, ih, iw)
53
+ out = net(data)
54
+
55
+ _, _, oh, ow = out.shape
56
+ print(f"Conv2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
57
+ macs, params = profile(net, inputs=(data,))
58
+ assert macs == n * out_c * oh * ow // g * in_c * kh * kw, (
59
+ f"{macs} v.s. {n * out_c * oh * ow // g * in_c * kh * kw}"
60
+ )
61
+
62
+ def test_convtranspose2d_no_bias(self):
63
+ """Tests a 2D ConvTranspose layer without bias using THOP profiling on predefined input dimensions and
64
+ parameters.
65
+ """
66
+ n, in_c, ih, iw = 1, 3, 2, 2
67
+ out_c, kh, kw = 1, 2, 2
68
+ s, p, d, g = 2, 0, 1, 1
69
+
70
+ net = nn.ConvTranspose2d(
71
+ in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False
72
+ )
73
+ data = torch.randn(n, in_c, ih, iw)
74
+ out = net(data)
75
+
76
+ _, _, oh, ow = out.shape
77
+
78
+ profile_result = profile(net, inputs=(data,))
79
+ macs = profile_result[0]
80
+ profile_result[1]
81
+ # For ConvTranspose: MACs = input_size * (output_channels / groups) * kernel_size
82
+ print(f"ConvTranspose2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
83
+ expected_macs = n * in_c * ih * iw * (out_c // g) * kh * kw
84
+ assert macs == expected_macs, f"{macs} v.s. {expected_macs}"
85
+
86
+ def test_convtranspose2d(self):
87
+ """Tests ConvTranspose2D layer with bias, profiling MACs and params for specific input dimensions and layer
88
+ configs.
89
+ """
90
+ n, in_c, ih, iw = 1, 3, 2, 2
91
+ out_c, kh, kw = 1, 2, 2
92
+ s, p, d, g = 2, 0, 1, 1
93
+
94
+ net = nn.ConvTranspose2d(
95
+ in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=True
96
+ )
97
+ data = torch.randn(n, in_c, ih, iw)
98
+ out = net(data)
99
+
100
+ _, _, oh, ow = out.shape
101
+
102
+ profile_result = profile(net, inputs=(data,))
103
+ macs = profile_result[0]
104
+ profile_result[1]
105
+ # For ConvTranspose: MACs = input_size * (output_channels / groups) * kernel_size
106
+ print(f"ConvTranspose2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
107
+ expected_macs = n * in_c * ih * iw * (out_c // g) * kh * kw
108
+ assert macs == expected_macs, f"{macs} v.s. {expected_macs}"
109
+
110
+ def test_convtranspose2d_groups(self):
111
+ """Tests ConvTranspose2D layer with groups, validating MAC calculation for grouped transposed convolution."""
112
+ n, in_c, ih, iw = 1, 8, 4, 4
113
+ out_c, kh, kw = 4, 3, 3
114
+ s, p, d, g = 1, 1, 1, 2
115
+
116
+ net = nn.ConvTranspose2d(
117
+ in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False
118
+ )
119
+ data = torch.randn(n, in_c, ih, iw)
120
+ out = net(data)
121
+
122
+ _, _, oh, ow = out.shape
123
+
124
+ profile_result = profile(net, inputs=(data,))
125
+ macs = profile_result[0]
126
+ profile_result[1]
127
+ # For ConvTranspose with groups: MACs = input_size * (output_channels / groups) * kernel_size
128
+ print(f"ConvTranspose2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={s}, padding={p}, out={oh}x{ow}")
129
+ expected_macs = n * in_c * ih * iw * (out_c // g) * kh * kw
130
+ assert macs == expected_macs, f"{macs} v.s. {expected_macs}"
131
+
132
+ def test_convtranspose2d_random(self):
133
+ """Test ConvTranspose2D layer with random parameters and validate the computed MACs and parameters using
134
+ 'profile'.
135
+ """
136
+ for _ in range(10):
137
+ # Generate random parameters ensuring valid ConvTranspose configurations
138
+ out_c, kh, kw = torch.randint(1, 10, (3,)).tolist()
139
+ n, in_c = torch.randint(1, 5, (2,)).tolist()
140
+ stride = int(torch.randint(1, 3, (1,)).item()) # stride
141
+ padding = int(torch.randint(0, 2, (1,)).item()) # padding
142
+ dilation, groups = 1, 1 # Keep dilation=1 and groups=1 for simplicity
143
+
144
+ # Ensure input size is large enough to produce valid output
145
+ # ConvTranspose output size formula: (input_size - 1) * stride - 2 * padding + kernel_size
146
+ # To ensure positive output, we need: input_size >= (2 * padding + 1) / stride + 1
147
+ min_input_size = max(3, (2 * padding + kh) // stride + 1, (2 * padding + kw) // stride + 1)
148
+ ih, iw = torch.randint(min_input_size, min_input_size + 10, (2,)).tolist()
149
+
150
+ net = nn.ConvTranspose2d(
151
+ in_c,
152
+ out_c,
153
+ kernel_size=(kh, kw),
154
+ stride=stride,
155
+ padding=padding,
156
+ dilation=dilation,
157
+ groups=groups,
158
+ bias=False,
159
+ )
160
+ data = torch.randn(n, in_c, ih, iw)
161
+ out = net(data)
162
+
163
+ _, _, oh, ow = out.shape
164
+
165
+ profile_result = profile(net, inputs=(data,))
166
+ macs = profile_result[0]
167
+ profile_result[1]
168
+ # For ConvTranspose: MACs = input_size * (output_channels / groups) * kernel_size
169
+ expected_macs = n * in_c * ih * iw * (out_c // groups) * kh * kw
170
+ print(
171
+ f"ConvTranspose2d: in={ih}x{iw}, kernel={kh}x{kw}, stride={stride}, padding={padding}, out={oh}x{ow}, macs={macs}"
172
+ )
173
+ assert macs == expected_macs, f"ConvTranspose2d MACs: {macs} v.s. {expected_macs}"
174
+
175
+ def test_conv_vs_convtranspose_symmetry(self):
176
+ """
177
+ Test that Conv2d and ConvTranspose2d with symmetric configurations have equal MAC counts.
178
+
179
+ Test case: Conv2d downsamples 4x4 -> 2x2, ConvTranspose2d upsamples 2x2 -> 4x4.
180
+ """
181
+ # Conv2d: 4x4 -> 2x2
182
+ conv_net = nn.Conv2d(1, 3, kernel_size=2, stride=2, bias=False)
183
+ conv_data = torch.randn(1, 1, 4, 4)
184
+ conv_out = conv_net(conv_data)
185
+ conv_profile_result = profile(conv_net, inputs=(conv_data,))
186
+ conv_macs = conv_profile_result[0]
187
+ conv_params = conv_profile_result[1]
188
+
189
+ # ConvTranspose2d: 2x2 -> 4x4 (symmetric operation)
190
+ convt_net = nn.ConvTranspose2d(3, 1, kernel_size=2, stride=2, bias=False)
191
+ convt_data = torch.randn(1, 3, 2, 2)
192
+ convt_out = convt_net(convt_data)
193
+ convt_profile_result = profile(convt_net, inputs=(convt_data,))
194
+ convt_macs = convt_profile_result[0]
195
+ convt_params = convt_profile_result[1]
196
+
197
+ # Verify symmetric operations have equal MAC counts
198
+ assert conv_macs == convt_macs, f"Symmetric operations should have equal MACs: {conv_macs} != {convt_macs}"
199
+
200
+ # Manual verification
201
+ # Conv: output_size * (input_channels / groups) * kernel_size
202
+ conv_expected = (
203
+ conv_out.numel()
204
+ * (conv_net.in_channels // conv_net.groups)
205
+ * (conv_net.kernel_size[0] * conv_net.kernel_size[1])
206
+ )
207
+ # ConvTranspose: input_size * (output_channels / groups) * kernel_size
208
+ convt_expected = (
209
+ convt_data.numel()
210
+ * (convt_net.out_channels // convt_net.groups)
211
+ * (convt_net.kernel_size[0] * convt_net.kernel_size[1])
212
+ )
213
+ print(f"Conv2d: {conv_data.shape} -> {conv_out.shape}, MACs: {conv_macs}, Params: {conv_params}")
214
+ print(f"ConvTranspose2d: {convt_data.shape} -> {convt_out.shape}, MACs: {convt_macs}, Params: {convt_params}")
215
+ print(f"Conv2d expected MACs: {conv_expected}, ConvTranspose2d expected MACs: {convt_expected}")
216
+ assert conv_macs == conv_expected, f"Conv2d MAC calculation incorrect: {conv_macs} != {conv_expected}"
217
+ assert convt_macs == convt_expected, (
218
+ f"ConvTranspose2d MAC calculation incorrect: {convt_macs} != {convt_expected}"
219
+ )
@@ -17,7 +17,7 @@ class TestUtils:
17
17
  print(flops, params)
18
18
  assert flops == n * in_c * out_c
19
19
 
20
- def test_matmul_case2(self):
20
+ def test_matmul_case3(self): # Note renamed to case3 by Glenn Jocher as duplicated above function name
21
21
  """Tests matrix multiplication to profile FLOPs and parameters of nn.Linear layer using random dimensions."""
22
22
  for _ in range(10):
23
23
  n, in_c, out_c = torch.randint(1, 500, (3,)).tolist()
@@ -1,7 +1,6 @@
1
1
  # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
2
 
3
- __version__ = "2.0.15"
4
-
3
+ __version__ = "2.0.17"
5
4
 
6
5
  import torch
7
6
 
@@ -9,3 +8,5 @@ from .profile import profile, profile_origin
9
8
  from .utils import clever_format
10
9
 
11
10
  default_dtype = torch.float64
11
+
12
+ __all__ = ["profile", "profile_origin", "clever_format", "default_dtype"]
@@ -6,6 +6,11 @@ from distutils.version import LooseVersion
6
6
  import torch
7
7
  import torch as th
8
8
  import torch.nn as nn
9
+ from torch.fx import symbolic_trace
10
+ from torch.fx.passes.shape_prop import ShapeProp
11
+
12
+ from .utils import prRed, prYellow
13
+ from .vision.calc_func import calculate_conv
9
14
 
10
15
  if LooseVersion(torch.__version__) < LooseVersion("1.8.0"):
11
16
  logging.warning(
@@ -40,9 +45,6 @@ def count_fn_linear(input_shapes, output_shapes, *args, **kwargs):
40
45
  return flops
41
46
 
42
47
 
43
- from .vision.calc_func import calculate_conv
44
-
45
-
46
48
  def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs):
47
49
  """Calculates total operations (FLOPs) for a 2D conv layer based on input and output shapes using
48
50
  `calculate_conv`.
@@ -118,11 +120,6 @@ for k in zero_ops:
118
120
 
119
121
  missing_maps = {}
120
122
 
121
- from torch.fx import symbolic_trace
122
- from torch.fx.passes.shape_prop import ShapeProp
123
-
124
- from .utils import prRed, prYellow
125
-
126
123
 
127
124
  def null_print(*args, **kwargs):
128
125
  """A no-op print function that takes any arguments without performing any actions."""
@@ -211,11 +208,15 @@ def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False):
211
208
  if __name__ == "__main__":
212
209
 
213
210
  class MyOP(nn.Module):
211
+ """Custom operator that performs a simple forward pass dividing input by 1."""
212
+
214
213
  def forward(self, input):
215
214
  """Performs forward pass on given input data."""
216
215
  return input / 1
217
216
 
218
217
  class MyModule(torch.nn.Module):
218
+ """Neural network module with two linear layers and a custom MyOP operator."""
219
+
219
220
  def __init__(self):
220
221
  """Initializes MyModule with two linear layers and a custom MyOP operator."""
221
222
  super().__init__()
@@ -1,7 +1,30 @@
1
1
  # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
2
 
3
- from thop.rnn_hooks import *
4
- from thop.vision.basic_hooks import *
3
+ from thop.rnn_hooks import (
4
+ count_gru,
5
+ count_gru_cell,
6
+ count_lstm,
7
+ count_lstm_cell,
8
+ count_rnn,
9
+ count_rnn_cell,
10
+ torch,
11
+ )
12
+ from thop.vision.basic_hooks import (
13
+ count_adap_avgpool,
14
+ count_avgpool,
15
+ count_convNd,
16
+ count_convtNd,
17
+ count_linear,
18
+ count_normalization,
19
+ count_parameters,
20
+ count_prelu,
21
+ count_relu,
22
+ count_softmax,
23
+ count_upsample,
24
+ logging,
25
+ nn,
26
+ zero_ops,
27
+ )
5
28
 
6
29
  from .utils import prRed
7
30
 
@@ -12,9 +35,9 @@ register_hooks = {
12
35
  nn.Conv1d: count_convNd,
13
36
  nn.Conv2d: count_convNd,
14
37
  nn.Conv3d: count_convNd,
15
- nn.ConvTranspose1d: count_convNd,
16
- nn.ConvTranspose2d: count_convNd,
17
- nn.ConvTranspose3d: count_convNd,
38
+ nn.ConvTranspose1d: count_convtNd,
39
+ nn.ConvTranspose2d: count_convtNd,
40
+ nn.ConvTranspose3d: count_convtNd,
18
41
  nn.BatchNorm1d: count_normalization,
19
42
  nn.BatchNorm2d: count_normalization,
20
43
  nn.BatchNorm3d: count_normalization,
@@ -2,10 +2,24 @@
2
2
 
3
3
  import logging
4
4
 
5
+ import torch
5
6
  import torch.nn as nn
6
7
  from torch.nn.modules.conv import _ConvNd
7
8
 
8
- from .calc_func import *
9
+ from thop.vision.calc_func import (
10
+ calculate_adaptive_avg,
11
+ calculate_avgpool,
12
+ calculate_conv,
13
+ calculate_conv2d_flops,
14
+ calculate_linear,
15
+ calculate_norm,
16
+ calculate_parameters,
17
+ calculate_relu,
18
+ calculate_relu_flops,
19
+ calculate_softmax,
20
+ calculate_upsample,
21
+ calculate_zero_ops,
22
+ )
9
23
 
10
24
  multiply_adds = 1
11
25
 
@@ -30,6 +44,7 @@ def count_convNd(m: _ConvNd, x, y: torch.Tensor):
30
44
  kernel_size=list(m.weight.shape),
31
45
  groups=m.groups,
32
46
  bias=m.bias,
47
+ transpose=False,
33
48
  )
34
49
  # N x Cout x H x W x (Cin x Kw x Kh + bias)
35
50
  # m.total_ops += calculate_conv(
@@ -41,6 +56,20 @@ def count_convNd(m: _ConvNd, x, y: torch.Tensor):
41
56
  # )
42
57
 
43
58
 
59
+ def count_convtNd(m: _ConvNd, x, y: torch.Tensor):
60
+ """Calculate and add the number of convolutional operations (FLOPs) for a ConvNd layer to the model's total ops."""
61
+ x = x[0]
62
+
63
+ m.total_ops += calculate_conv2d_flops(
64
+ input_size=list(x.shape),
65
+ output_size=list(y.shape),
66
+ kernel_size=list(m.weight.shape),
67
+ groups=m.groups,
68
+ bias=m.bias,
69
+ transpose=True,
70
+ )
71
+
72
+
44
73
  def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor):
45
74
  """Calculates and updates total operations (FLOPs) for a convolutional layer in a PyTorch model."""
46
75
  x = x[0]
@@ -29,13 +29,18 @@ def calculate_zero_ops():
29
29
  return torch.DoubleTensor([0])
30
30
 
31
31
 
32
- def calculate_conv2d_flops(input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False):
32
+ def calculate_conv2d_flops(
33
+ input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False, transpose: bool = False
34
+ ):
33
35
  """Calculate FLOPs for a Conv2D layer using input/output sizes, kernel size, groups, and the bias flag."""
34
36
  # n, in_c, ih, iw = input_size
35
37
  # 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:])
38
+ if transpose:
39
+ out_c = output_size[1]
40
+ return l_prod(input_size) * (out_c // groups) * l_prod(kernel_size[2:])
41
+ else:
42
+ in_c = input_size[1]
43
+ return l_prod(output_size) * (in_c // groups) * l_prod(kernel_size[2:])
39
44
 
40
45
 
41
46
  def calculate_conv(bias, kernel_size, output_size, in_channel, group):
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: ultralytics-thop
3
- Version: 2.0.15
3
+ Version: 2.0.17
4
4
  Summary: Ultralytics THOP package for fast computation of PyTorch model FLOPs and parameters.
5
5
  Author-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>
6
6
  Maintainer-email: Ultralytics <hello@ultralytics.com>
@@ -56,7 +56,7 @@ THOP offers an intuitive API designed to profile PyTorch models by calculating t
56
56
 
57
57
  Get started with THOP quickly by installing it via pip:
58
58
 
59
- [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://www.pepy.tech/projects/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
59
+ [![PyPI - Version](https://img.shields.io/pypi/v/ultralytics-thop?logo=pypi&logoColor=white)](https://pypi.org/project/ultralytics-thop/) [![Downloads](https://static.pepy.tech/badge/ultralytics-thop)](https://clickpy.clickhouse.com/dashboard/ultralytics-thop) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/ultralytics-thop?logo=python&logoColor=gold)](https://pypi.org/project/ultralytics-thop/)
60
60
 
61
61
  ```bash
62
62
  pip install ultralytics-thop
@@ -74,7 +74,7 @@ This ensures you have the most recent version, incorporating the latest improvem
74
74
 
75
75
  ### Basic Usage
76
76
 
77
- Profiling a standard PyTorch model like [ResNet50](https://pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
77
+ Profiling a standard PyTorch model like [ResNet50](https://docs.pytorch.org/vision/main/models/generated/torchvision.models.resnet50.html) is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the `profile` function:
78
78
 
79
79
  ```python
80
80
  import torch
@@ -230,7 +230,7 @@ The table below showcases the parameters and MACs for several popular [computer
230
230
 
231
231
  ## 🙌 Contribute
232
232
 
233
- We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
233
+ We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new [PyTorch layers](https://docs.pytorch.org/docs/stable/nn.html), improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our [Contributing Guide](https://docs.ultralytics.com/help/contributing/) for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the [machine learning](https://www.ultralytics.com/glossary/machine-learning-ml) community. Don't hesitate to share your feedback and suggestions!
234
234
 
235
235
  ## 📜 License
236
236
 
@@ -1,61 +0,0 @@
1
- # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
2
-
3
- import torch
4
- import torch.nn as nn
5
-
6
- from thop import profile
7
-
8
-
9
- class TestUtils:
10
- """Utility class for testing Conv2D layers with and without bias, profiling FLOPs and parameters using THOP."""
11
-
12
- def test_conv2d_no_bias(self):
13
- """Tests a 2D Conv layer without bias using THOP profiling on predefined input dimensions and parameters."""
14
- n, in_c, ih, iw = 1, 3, 32, 32 # torch.randint(1, 10, (4,)).tolist()
15
- out_c, kh, kw = 12, 5, 5
16
- s, p, d, g = 1, 1, 1, 1
17
-
18
- net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False)
19
- data = torch.randn(n, in_c, ih, iw)
20
- out = net(data)
21
-
22
- _, _, oh, ow = out.shape
23
-
24
- flops, params = profile(net, inputs=(data,))
25
- assert flops == 810000, f"{flops} v.s. 810000"
26
-
27
- def test_conv2d(self):
28
- """Tests Conv2D layer with bias, profiling FLOPs and params for specific input dimensions and layer configs."""
29
- n, in_c, ih, iw = 1, 3, 32, 32 # torch.randint(1, 10, (4,)).tolist()
30
- out_c, kh, kw = 12, 5, 5
31
- s, p, d, g = 1, 1, 1, 1
32
-
33
- net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=True)
34
- data = torch.randn(n, in_c, ih, iw)
35
- out = net(data)
36
-
37
- _, _, oh, ow = out.shape
38
-
39
- flops, params = profile(net, inputs=(data,))
40
- assert flops == 810000, f"{flops} v.s. 810000"
41
-
42
- def test_conv2d_random(self):
43
- """Test Conv2D layer with random parameters and validate the computed FLOPs and parameters using 'profile'."""
44
- for _ in range(10):
45
- out_c, kh, kw = torch.randint(1, 20, (3,)).tolist()
46
- n, in_c, ih, iw = torch.randint(1, 20, (4,)).tolist() # torch.randint(1, 10, (4,)).tolist()
47
- ih += kh
48
- iw += kw
49
- s, p, d, g = 1, 1, 1, 1
50
-
51
- net = nn.Conv2d(in_c, out_c, kernel_size=(kh, kw), stride=s, padding=p, dilation=d, groups=g, bias=False)
52
- data = torch.randn(n, in_c, ih, iw)
53
- out = net(data)
54
-
55
- _, _, oh, ow = out.shape
56
-
57
- flops, params = profile(net, inputs=(data,))
58
- print(flops, params)
59
- assert flops == n * out_c * oh * ow // g * in_c * kh * kw, (
60
- f"{flops} v.s. {n * out_c * oh * ow // g * in_c * kh * kw}"
61
- )