ultralytics-thop 0.2.4__tar.gz → 0.2.5__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 (25) hide show
  1. {ultralytics_thop-0.2.4/ultralytics_thop.egg-info → ultralytics_thop-0.2.5}/PKG-INFO +1 -1
  2. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/tests/test_conv2d.py +3 -5
  3. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/tests/test_matmul.py +0 -1
  4. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/tests/test_relu.py +0 -1
  5. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/tests/test_utils.py +0 -2
  6. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/__init__.py +1 -1
  7. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/fx_profile.py +3 -5
  8. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/profile.py +4 -6
  9. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/rnn_hooks.py +33 -36
  10. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/utils.py +1 -12
  11. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/vision/basic_hooks.py +1 -3
  12. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/vision/calc_func.py +7 -15
  13. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5/ultralytics_thop.egg-info}/PKG-INFO +1 -1
  14. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/ultralytics_thop.egg-info/SOURCES.txt +0 -3
  15. ultralytics_thop-0.2.4/thop/onnx_profile.py +0 -82
  16. ultralytics_thop-0.2.4/thop/vision/efficientnet.py +0 -9
  17. ultralytics_thop-0.2.4/thop/vision/onnx_counter.py +0 -371
  18. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/LICENSE +0 -0
  19. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/README.md +0 -0
  20. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/pyproject.toml +0 -0
  21. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/setup.cfg +0 -0
  22. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/thop/vision/__init__.py +0 -0
  23. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/ultralytics_thop.egg-info/dependency_links.txt +0 -0
  24. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/ultralytics_thop.egg-info/requires.txt +0 -0
  25. {ultralytics_thop-0.2.4 → ultralytics_thop-0.2.5}/ultralytics_thop.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ultralytics-thop
3
- Version: 0.2.4
3
+ Version: 0.2.5
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: Glenn Jocher
@@ -1,7 +1,5 @@
1
- import pytest
2
1
  import torch
3
2
  import torch.nn as nn
4
- from jinja2 import StrictUndefined
5
3
 
6
4
  from thop import profile
7
5
 
@@ -22,7 +20,7 @@ class TestUtils:
22
20
  _, _, oh, ow = out.shape
23
21
 
24
22
  flops, params = profile(net, inputs=(data,))
25
- assert flops == 810000, f"{flops} v.s. {810000}"
23
+ assert flops == 810000, f"{flops} v.s. 810000"
26
24
 
27
25
  def test_conv2d(self):
28
26
  """Tests a Conv2D layer with specific input dimensions, kernel size, stride, padding, dilation, and groups."""
@@ -37,11 +35,11 @@ class TestUtils:
37
35
  _, _, oh, ow = out.shape
38
36
 
39
37
  flops, params = profile(net, inputs=(data,))
40
- assert flops == 810000, f"{flops} v.s. {810000}"
38
+ assert flops == 810000, f"{flops} v.s. 810000"
41
39
 
42
40
  def test_conv2d_random(self):
43
41
  """Test Conv2D layer with random parameters and validate the computed FLOPs and parameters using 'profile'."""
44
- for i in range(10):
42
+ for _ in range(10):
45
43
  out_c, kh, kw = torch.randint(1, 20, (3,)).tolist()
46
44
  n, in_c, ih, iw = torch.randint(1, 20, (4,)).tolist() # torch.randint(1, 10, (4,)).tolist()
47
45
  ih += kh
@@ -1,4 +1,3 @@
1
- import pytest
2
1
  import torch
3
2
  import torch.nn as nn
4
3
 
@@ -1,4 +1,3 @@
1
- import pytest
2
1
  import torch
3
2
  import torch.nn as nn
4
3
 
@@ -1,5 +1,3 @@
1
- import pytest
2
-
3
1
  from thop import utils
4
2
 
5
3
 
@@ -1,4 +1,4 @@
1
- __version__ = "0.2.4"
1
+ __version__ = "0.2.5"
2
2
 
3
3
  import torch
4
4
 
@@ -85,9 +85,7 @@ def count_nn_bn2d(module: nn.BatchNorm2d, input_shapes, output_shapes):
85
85
  """Calculate the total operations for a given nn.BatchNorm2d module based on its output shape."""
86
86
  assert len(output_shapes) == 1, "nn.BatchNorm2d should only have one output"
87
87
  y = output_shapes[0]
88
- # y = (x - mean) / \sqrt{var + e} * weight + bias
89
- total_ops = 2 * y.numel()
90
- return total_ops
88
+ return 2 * y.numel()
91
89
 
92
90
 
93
91
  zero_ops = (
@@ -120,7 +118,7 @@ missing_maps = {}
120
118
  from torch.fx import symbolic_trace
121
119
  from torch.fx.passes.shape_prop import ShapeProp
122
120
 
123
- from .utils import prGreen, prRed, prYellow
121
+ from .utils import prRed, prYellow
124
122
 
125
123
 
126
124
  def null_print(*args, **kwargs):
@@ -193,7 +191,7 @@ def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False):
193
191
  prRed(f"{key} is missing")
194
192
  print("module type:", type(m))
195
193
  if isinstance(m, zero_ops):
196
- print(f"weight_shape: None")
194
+ print("weight_shape: None")
197
195
  else:
198
196
  print(type(m))
199
197
  print(f"weight_shape: {mod.state_dict()[node.target + '.weight'].shape}")
@@ -1,9 +1,7 @@
1
1
  from thop.rnn_hooks import *
2
2
  from thop.vision.basic_hooks import *
3
3
 
4
- # logger = logging.getLogger(__name__)
5
- # logger.setLevel(logging.INFO)
6
- from .utils import prGreen, prRed, prYellow
4
+ from .utils import prRed
7
5
 
8
6
  default_dtype = torch.float64
9
7
 
@@ -68,7 +66,7 @@ def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=
68
66
  verbose = True
69
67
 
70
68
  def add_hooks(m):
71
- if len(list(m.children())) > 0:
69
+ if list(m.children()):
72
70
  return
73
71
 
74
72
  if hasattr(m, "total_ops") or hasattr(m, "total_params"):
@@ -114,7 +112,7 @@ def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=
114
112
  total_ops = 0
115
113
  total_params = 0
116
114
  for m in model.modules():
117
- if len(list(m.children())) > 0: # skip for non-leaf module
115
+ if list(m.children()): # skip for non-leaf module
118
116
  continue
119
117
  total_ops += m.total_ops
120
118
  total_params += m.total_params
@@ -129,7 +127,7 @@ def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=
129
127
 
130
128
  # remove temporal buffers
131
129
  for n, m in model.named_modules():
132
- if len(list(m.children())) > 0:
130
+ if list(m.children()):
133
131
  continue
134
132
  if "total_ops" in m._buffers:
135
133
  m._buffers.pop("total_ops")
@@ -103,13 +103,12 @@ def count_rnn(m: nn.RNN, x, y):
103
103
  if isinstance(x[0], PackedSequence):
104
104
  batch_size = torch.max(x[0].batch_sizes)
105
105
  num_steps = x[0].batch_sizes.size(0)
106
+ elif m.batch_first:
107
+ batch_size = x[0].size(0)
108
+ num_steps = x[0].size(1)
106
109
  else:
107
- if m.batch_first:
108
- batch_size = x[0].size(0)
109
- num_steps = x[0].size(1)
110
- else:
111
- batch_size = x[0].size(1)
112
- num_steps = x[0].size(0)
110
+ batch_size = x[0].size(1)
111
+ num_steps = x[0].size(0)
113
112
 
114
113
  total_ops = 0
115
114
  if m.bidirectional:
@@ -117,12 +116,12 @@ def count_rnn(m: nn.RNN, x, y):
117
116
  else:
118
117
  total_ops += _count_rnn_cell(input_size, hidden_size, bias)
119
118
 
120
- for i in range(num_layers - 1):
121
- if m.bidirectional:
122
- total_ops += _count_rnn_cell(hidden_size * 2, hidden_size, bias) * 2
123
- else:
124
- total_ops += _count_rnn_cell(hidden_size, hidden_size, bias)
125
-
119
+ for _ in range(num_layers - 1):
120
+ total_ops += (
121
+ _count_rnn_cell(hidden_size * 2, hidden_size, bias) * 2
122
+ if m.bidirectional
123
+ else _count_rnn_cell(hidden_size, hidden_size, bias)
124
+ )
126
125
  # time unroll
127
126
  total_ops *= num_steps
128
127
  # batch_size
@@ -141,13 +140,12 @@ def count_gru(m: nn.GRU, x, y):
141
140
  if isinstance(x[0], PackedSequence):
142
141
  batch_size = torch.max(x[0].batch_sizes)
143
142
  num_steps = x[0].batch_sizes.size(0)
143
+ elif m.batch_first:
144
+ batch_size = x[0].size(0)
145
+ num_steps = x[0].size(1)
144
146
  else:
145
- if m.batch_first:
146
- batch_size = x[0].size(0)
147
- num_steps = x[0].size(1)
148
- else:
149
- batch_size = x[0].size(1)
150
- num_steps = x[0].size(0)
147
+ batch_size = x[0].size(1)
148
+ num_steps = x[0].size(0)
151
149
 
152
150
  total_ops = 0
153
151
  if m.bidirectional:
@@ -155,12 +153,12 @@ def count_gru(m: nn.GRU, x, y):
155
153
  else:
156
154
  total_ops += _count_gru_cell(input_size, hidden_size, bias)
157
155
 
158
- for i in range(num_layers - 1):
159
- if m.bidirectional:
160
- total_ops += _count_gru_cell(hidden_size * 2, hidden_size, bias) * 2
161
- else:
162
- total_ops += _count_gru_cell(hidden_size, hidden_size, bias)
163
-
156
+ for _ in range(num_layers - 1):
157
+ total_ops += (
158
+ _count_gru_cell(hidden_size * 2, hidden_size, bias) * 2
159
+ if m.bidirectional
160
+ else _count_gru_cell(hidden_size, hidden_size, bias)
161
+ )
164
162
  # time unroll
165
163
  total_ops *= num_steps
166
164
  # batch_size
@@ -181,13 +179,12 @@ def count_lstm(m: nn.LSTM, x, y):
181
179
  if isinstance(x[0], PackedSequence):
182
180
  batch_size = torch.max(x[0].batch_sizes)
183
181
  num_steps = x[0].batch_sizes.size(0)
182
+ elif m.batch_first:
183
+ batch_size = x[0].size(0)
184
+ num_steps = x[0].size(1)
184
185
  else:
185
- if m.batch_first:
186
- batch_size = x[0].size(0)
187
- num_steps = x[0].size(1)
188
- else:
189
- batch_size = x[0].size(1)
190
- num_steps = x[0].size(0)
186
+ batch_size = x[0].size(1)
187
+ num_steps = x[0].size(0)
191
188
 
192
189
  total_ops = 0
193
190
  if m.bidirectional:
@@ -195,12 +192,12 @@ def count_lstm(m: nn.LSTM, x, y):
195
192
  else:
196
193
  total_ops += _count_lstm_cell(input_size, hidden_size, bias)
197
194
 
198
- for i in range(num_layers - 1):
199
- if m.bidirectional:
200
- total_ops += _count_lstm_cell(hidden_size * 2, hidden_size, bias) * 2
201
- else:
202
- total_ops += _count_lstm_cell(hidden_size, hidden_size, bias)
203
-
195
+ for _ in range(num_layers - 1):
196
+ total_ops += (
197
+ _count_lstm_cell(hidden_size * 2, hidden_size, bias) * 2
198
+ if m.bidirectional
199
+ else _count_lstm_cell(hidden_size, hidden_size, bias)
200
+ )
204
201
  # time unroll
205
202
  total_ops *= num_steps
206
203
  # batch_size
@@ -20,15 +20,6 @@ prRed = colorful_print(print, color=COLOR_RED)
20
20
  prGreen = colorful_print(print, color=COLOR_GREEN)
21
21
  prYellow = colorful_print(print, color=COLOR_YELLOW)
22
22
 
23
- # def prRed(skk):
24
- # print("\033[91m{}\033[00m".format(skk))
25
-
26
- # def prGreen(skk):
27
- # print("\033[92m{}\033[00m".format(skk))
28
-
29
- # def prYellow(skk):
30
- # print("\033[93m{}\033[00m".format(skk))
31
-
32
23
 
33
24
  def clever_format(nums, format="%.2f"):
34
25
  """Formats numerical values into a more readable string with units (K, M, G, T) based on their magnitude."""
@@ -48,9 +39,7 @@ def clever_format(nums, format="%.2f"):
48
39
  else:
49
40
  clever_nums.append(format % num + "B")
50
41
 
51
- clever_nums = clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
52
-
53
- return clever_nums
42
+ return clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
54
43
 
55
44
 
56
45
  if __name__ == "__main__":
@@ -1,7 +1,5 @@
1
- import argparse
2
1
  import logging
3
2
 
4
- import torch
5
3
  import torch.nn as nn
6
4
  from torch.nn.modules.conv import _ConvNd
7
5
 
@@ -139,7 +137,7 @@ def count_upsample(m, x, y):
139
137
  "bilinear",
140
138
  "bicubic",
141
139
  ): # "trilinear"
142
- logging.warning("mode %s is not implemented yet, take it a zero op" % m.mode)
140
+ logging.warning(f"mode {m.mode} is not implemented yet, take it a zero op")
143
141
  m.total_ops += 0
144
142
  else:
145
143
  x = x[0]
@@ -14,23 +14,17 @@ def l_prod(in_list):
14
14
 
15
15
  def l_sum(in_list):
16
16
  """Calculate the sum of all elements in a list."""
17
- res = 0
18
- for _ in in_list:
19
- res += _
20
- return res
17
+ return sum(in_list)
21
18
 
22
19
 
23
20
  def calculate_parameters(param_list):
24
21
  """Calculate the total number of parameters in a list of tensors."""
25
- total_params = 0
26
- for p in param_list:
27
- total_params += torch.DoubleTensor([p.nelement()])
28
- return total_params
22
+ return sum(torch.DoubleTensor([p.nelement()]) for p in param_list)
29
23
 
30
24
 
31
25
  def calculate_zero_ops():
32
26
  """Return a tensor initialized to zero."""
33
- return torch.DoubleTensor([int(0)])
27
+ return torch.DoubleTensor([0])
34
28
 
35
29
 
36
30
  def calculate_conv2d_flops(input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False):
@@ -90,14 +84,12 @@ def calculate_adaptive_avg(kernel_size, output_size):
90
84
  def calculate_upsample(mode: str, output_size):
91
85
  """Calculate the number of operations for upsample methods given the mode and output size."""
92
86
  total_ops = output_size
93
- if mode == "linear":
94
- total_ops *= 5
87
+ if mode == "bicubic":
88
+ total_ops *= 224 + 35
95
89
  elif mode == "bilinear":
96
90
  total_ops *= 11
97
- elif mode == "bicubic":
98
- ops_solve_A = 224 # 128 muls + 96 adds
99
- ops_solve_p = 35 # 16 muls + 12 adds + 4 muls + 3 adds
100
- total_ops *= ops_solve_A + ops_solve_p
91
+ elif mode == "linear":
92
+ total_ops *= 5
101
93
  elif mode == "trilinear":
102
94
  total_ops *= 13 * 2 + 5
103
95
  return torch.DoubleTensor([int(total_ops)])
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ultralytics-thop
3
- Version: 0.2.4
3
+ Version: 0.2.5
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: Glenn Jocher
@@ -7,15 +7,12 @@ tests/test_relu.py
7
7
  tests/test_utils.py
8
8
  thop/__init__.py
9
9
  thop/fx_profile.py
10
- thop/onnx_profile.py
11
10
  thop/profile.py
12
11
  thop/rnn_hooks.py
13
12
  thop/utils.py
14
13
  thop/vision/__init__.py
15
14
  thop/vision/basic_hooks.py
16
15
  thop/vision/calc_func.py
17
- thop/vision/efficientnet.py
18
- thop/vision/onnx_counter.py
19
16
  ultralytics_thop.egg-info/PKG-INFO
20
17
  ultralytics_thop.egg-info/SOURCES.txt
21
18
  ultralytics_thop.egg-info/dependency_links.txt
@@ -1,82 +0,0 @@
1
- import numpy as np
2
- import onnx
3
- import torch
4
- import torch.nn
5
- from onnx import numpy_helper
6
-
7
- from thop.vision.onnx_counter import onnx_operators
8
-
9
-
10
- class OnnxProfile:
11
- def __init__(self):
12
- """Initialize the OnnxProfile class with necessary imports for ONNX profiling."""
13
- pass
14
-
15
- def calculate_params(self, model: onnx.ModelProto):
16
- """Calculate the total number of parameters in an ONNX model."""
17
- onnx_weights = model.graph.initializer
18
- params = 0
19
-
20
- for onnx_w in onnx_weights:
21
- try:
22
- weight = numpy_helper.to_array(onnx_w)
23
- params += np.prod(weight.shape)
24
- except Exception as _:
25
- pass
26
-
27
- return params
28
-
29
- def create_dict(self, weight, input, output):
30
- """Create and return a dictionary mapping weight, input, and output names to their respective dimensions."""
31
- diction = {}
32
- for w in weight:
33
- dim = np.array(w.dims)
34
- diction[str(w.name)] = dim
35
- if dim.size == 1:
36
- diction[str(w.name)] = np.append(1, dim)
37
- for i in input:
38
- # print(i.type.tensor_type.shape.dim[0].dim_value)
39
- dim = np.array(i.type.tensor_type.shape.dim[0].dim_value)
40
- # print(i.type.tensor_type.shape.dim.__sizeof__())
41
- # name2dims[str(i.name)] = [dim]
42
- dim = []
43
- for key in i.type.tensor_type.shape.dim:
44
- dim = np.append(dim, int(key.dim_value))
45
- # print(key.dim_value)
46
- # print(dim)
47
- diction[str(i.name)] = dim
48
- if dim.size == 1:
49
- diction[str(i.name)] = np.append(1, dim)
50
- for o in output:
51
- dim = np.array(o.type.tensor_type.shape.dim[0].dim_value)
52
- diction[str(o.name)] = [dim]
53
- if dim.size == 1:
54
- diction[str(o.name)] = np.append(1, dim)
55
- return diction
56
-
57
- def nodes_counter(self, diction, node):
58
- """Count nodes of a specific type in an ONNX graph, returning the count and associated node operation
59
- details.
60
- """
61
- if node.op_type not in onnx_operators:
62
- print("Sorry, we haven't add ", node.op_type, "into dictionary.")
63
- return 0, None, None
64
- else:
65
- fn = onnx_operators[node.op_type]
66
- return fn(diction, node)
67
-
68
- def calculate_macs(self, model: onnx.ModelProto) -> torch.DoubleTensor:
69
- macs = 0
70
- name2dims = {}
71
- weight = model.graph.initializer
72
- nodes = model.graph.node
73
- input = model.graph.input
74
- output = model.graph.output
75
- name2dims = self.create_dict(weight, input, output)
76
- macs = 0
77
- for n in nodes:
78
- macs_adding, out_size, outname = self.nodes_counter(name2dims, n)
79
-
80
- name2dims[outname] = out_size
81
- macs += macs_adding
82
- return np.array(macs[0])
@@ -1,9 +0,0 @@
1
- import argparse
2
- import logging
3
-
4
- import torch
5
- import torch.nn as nn
6
- from efficientnet_pytorch.utils import Conv2dDynamicSamePadding, Conv2dStaticSamePadding
7
- from torch.nn.modules.conv import _ConvNd
8
-
9
- register_hooks = {}
@@ -1,371 +0,0 @@
1
- import numpy as np
2
- import torch
3
- from onnx import numpy_helper
4
-
5
- from thop.vision.basic_hooks import zero_ops
6
-
7
- from .calc_func import (
8
- calculate_avgpool,
9
- calculate_conv,
10
- calculate_norm,
11
- calculate_softmax,
12
- calculate_zero_ops,
13
- counter_div,
14
- counter_matmul,
15
- counter_mul,
16
- counter_pow,
17
- counter_sqrt,
18
- )
19
-
20
-
21
- def onnx_counter_matmul(diction, node):
22
- """Calculates multiply-accumulate operations and output size for matrix multiplication in an ONNX model node."""
23
- input1 = node.input[0]
24
- input2 = node.input[1]
25
- input1_dim = diction[input1]
26
- input2_dim = diction[input2]
27
- out_size = np.append(input1_dim[0:-1], input2_dim[-1])
28
- output_name = node.output[0]
29
- macs = counter_matmul(input1_dim, out_size[-2:])
30
- return macs, out_size, output_name
31
-
32
-
33
- def onnx_counter_add(diction, node):
34
- """Calculate multiply-accumulate operations (MACs), output size, and output name for ONNX addition nodes."""
35
- if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size:
36
- out_size = diction[node.input[1]]
37
- else:
38
- out_size = diction[node.input[0]]
39
- output_name = node.output[0]
40
- macs = calculate_zero_ops()
41
- # if '140' in diction:
42
- # print(diction['140'],output_name)
43
- return macs, out_size, output_name
44
-
45
-
46
- def onnx_counter_conv(diction, node):
47
- """Calculates MACs, output size, and name for an ONNX convolution node based on input tensor dimensions and node
48
- attributes.
49
- """
50
- # bias,kernelsize,outputsize
51
- dim_bias = 0
52
- input_count = 0
53
- for i in node.input:
54
- input_count += 1
55
- if input_count == 3:
56
- dim_bias = 1
57
- dim_weight = diction[node.input[1]]
58
- else:
59
- dim_weight = diction[node.input[1]]
60
- for attr in node.attribute:
61
- # print(attr)
62
- if attr.name == "kernel_shape":
63
- dim_kernel = attr.ints # kw,kh
64
- if attr.name == "strides":
65
- dim_stride = attr.ints
66
- if attr.name == "pads":
67
- dim_pad = attr.ints
68
- if attr.name == "dilations":
69
- dim_dil = attr.ints
70
- if attr.name == "group":
71
- group = attr.i
72
- # print(dim_dil)
73
- dim_input = diction[node.input[0]]
74
- output_size = np.append(dim_input[0 : -np.array(dim_kernel).size - 1], dim_weight[0])
75
- hw = np.array(dim_input[-np.array(dim_kernel).size :])
76
- for i in range(hw.size):
77
- hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_dil[i] * (dim_kernel[i] - 1) - 1) / dim_stride[i] + 1)
78
- output_size = np.append(output_size, hw)
79
- macs = calculate_conv(dim_bias, np.prod(dim_kernel), np.prod(output_size), dim_weight[1], group)
80
- output_name = node.output[0]
81
-
82
- # if '140' in diction:
83
- # print("conv",diction['140'],output_name)
84
- return macs, output_size, output_name
85
-
86
-
87
- def onnx_counter_constant(diction, node):
88
- """Calculate MACs, output size, and output name for a constant operation in an ONNX model."""
89
- macs = calculate_zero_ops()
90
- output_name = node.output[0]
91
- output_size = [1]
92
- # print(macs, output_size, output_name)
93
- return macs, output_size, output_name
94
-
95
-
96
- def onnx_counter_mul(diction, node):
97
- """Calculate MACs, output size, and output name for a multiplication operation in an ONNX model."""
98
- if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size:
99
- input_size = diction[node.input[1]]
100
- else:
101
- input_size = diction[node.input[0]]
102
- macs = counter_mul(np.prod(input_size))
103
- output_size = diction[node.input[0]]
104
- output_name = node.output[0]
105
- return macs, output_size, output_name
106
-
107
-
108
- def onnx_counter_bn(diction, node):
109
- """Calculates MACs, output size, and output name for batch normalization layers in an ONNX model."""
110
- input_size = diction[node.input[0]]
111
- macs = calculate_norm(np.prod(input_size))
112
- output_name = node.output[0]
113
- output_size = input_size
114
- return macs, output_size, output_name
115
-
116
-
117
- def onnx_counter_relu(diction, node):
118
- """Calculates MACs, output size, and output name for ReLU layers in an ONNX model."""
119
- input_size = diction[node.input[0]]
120
- macs = calculate_zero_ops()
121
- output_name = node.output[0]
122
- output_size = input_size
123
- # print(macs, output_size, output_name)
124
- # if '140' in diction:
125
- # print("relu",diction['140'],output_name)
126
- return macs, output_size, output_name
127
-
128
-
129
- def onnx_counter_reducemean(diction, node):
130
- """Compute MACs, output size, and name for the ReduceMean ONNX node, adjusting dimensions based on the 'axes' and
131
- 'keepdims' attributes.
132
- """
133
- keep_dim = 0
134
- for attr in node.attribute:
135
- if "axes" in attr.name:
136
- dim_axis = np.array(attr.ints)
137
- elif "keepdims" in attr.name:
138
- keep_dim = attr.i
139
-
140
- input_size = diction[node.input[0]]
141
- macs = calculate_zero_ops()
142
- output_name = node.output[0]
143
- if keep_dim == 1:
144
- output_size = input_size
145
- else:
146
- output_size = np.delete(input_size, dim_axis)
147
- # output_size = input_size
148
- return macs, output_size, output_name
149
-
150
-
151
- def onnx_counter_sub(diction, node):
152
- """Computes MACs, output size, and output name for a given ONNX node with specified input size."""
153
- input_size = diction[node.input[0]]
154
- macs = calculate_zero_ops()
155
- output_name = node.output[0]
156
- output_size = input_size
157
- return macs, output_size, output_name
158
-
159
-
160
- def onnx_counter_pow(diction, node):
161
- """Calculates MACs, output size, and output name for a given ONNX 'Pow' node with specified input size."""
162
- if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size:
163
- input_size = diction[node.input[1]]
164
- else:
165
- input_size = diction[node.input[0]]
166
- macs = counter_pow(np.prod(input_size))
167
- output_name = node.output[0]
168
- output_size = input_size
169
- return macs, output_size, output_name
170
-
171
-
172
- def onnx_counter_sqrt(diction, node):
173
- """Calculate MACs and output information for the SQRT operation in an ONNX node."""
174
- input_size = diction[node.input[0]]
175
- macs = counter_sqrt(np.prod(input_size))
176
- output_name = node.output[0]
177
- output_size = input_size
178
- return macs, output_size, output_name
179
-
180
-
181
- def onnx_counter_div(diction, node):
182
- """Calculate MACs and output information for the DIV operation in an ONNX node."""
183
- if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size:
184
- input_size = diction[node.input[1]]
185
- else:
186
- input_size = diction[node.input[0]]
187
- macs = counter_div(np.prod(input_size))
188
- output_name = node.output[0]
189
- output_size = input_size
190
- return macs, output_size, output_name
191
-
192
-
193
- def onnx_counter_instance(diction, node):
194
- """Calculate MACs, output size, and name for an ONNX node instance."""
195
- input_size = diction[node.input[0]]
196
- macs = calculate_norm(np.prod(input_size))
197
- output_name = node.output[0]
198
- output_size = input_size
199
- return macs, output_size, output_name
200
-
201
-
202
- def onnx_counter_softmax(diction, node):
203
- """Calculate MACs, output size, and name for an ONNX softmax node instance."""
204
- input_size = diction[node.input[0]]
205
- dim = node.attribute[0].i
206
- nfeatures = input_size[dim]
207
- batch_size = np.prod(input_size) / nfeatures
208
- macs = calculate_softmax(nfeatures, batch_size)
209
- output_name = node.output[0]
210
- output_size = input_size
211
- return macs, output_size, output_name
212
-
213
-
214
- def onnx_counter_pad(diction, node):
215
- """Compute memory access cost (MACs), output size, and output name for ONNX pad operation."""
216
- # if
217
- # if (np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size):
218
- # input_size = diction[node.input[1]]
219
- # else:
220
- # input_size = diction[node.input[0]]
221
- input_size = diction[node.input[0]]
222
- macs = calculate_zero_ops()
223
- output_name = node.output[0]
224
- output_size = input_size
225
- return macs, output_size, output_name
226
-
227
-
228
- def onnx_counter_averagepool(diction, node):
229
- """Calculate MACs and output size for an AveragePool ONNX operation based on input dimensions and attributes."""
230
- macs = calculate_avgpool(np.prod(diction[node.input[0]]))
231
- output_name = node.output[0]
232
- dim_pad = None
233
- for attr in node.attribute:
234
- # print(attr)
235
- if attr.name == "kernel_shape":
236
- dim_kernel = attr.ints # kw,kh
237
- elif attr.name == "strides":
238
- dim_stride = attr.ints
239
- elif attr.name == "pads":
240
- dim_pad = attr.ints
241
- elif attr.name == "dilations":
242
- dim_dil = attr.ints
243
- # print(dim_dil)
244
- dim_input = diction[node.input[0]]
245
- hw = dim_input[-np.array(dim_kernel).size :]
246
- if dim_pad is not None:
247
- for i in range(hw.size):
248
- hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_kernel[i]) / dim_stride[i] + 1)
249
- output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw)
250
- else:
251
- for i in range(hw.size):
252
- hw[i] = int((hw[i] - dim_kernel[i]) / dim_stride[i] + 1)
253
- output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw)
254
- # print(macs, output_size, output_name)
255
- return macs, output_size, output_name
256
-
257
-
258
- def onnx_counter_flatten(diction, node):
259
- """Returns MACs, output size, and output name for an ONNX Flatten node."""
260
- macs = calculate_zero_ops()
261
- output_name = node.output[0]
262
- axis = node.attribute[0].i
263
- input_size = diction[node.input[0]]
264
- output_size = np.append(input_size[axis - 1], np.prod(input_size[axis:]))
265
- # print("flatten",output_size)
266
- return macs, output_size, output_name
267
-
268
-
269
- def onnx_counter_gemm(diction, node):
270
- """Calculate multiply–accumulate operations (MACs), output size, and name for ONNX Gemm node."""
271
- # Compute Y = alpha * A' * B' + beta * C
272
- input_size = diction[node.input[0]]
273
- dim_weight = diction[node.input[1]]
274
- # print(input_size,dim_weight)
275
- macs = np.prod(input_size) * dim_weight[1] + dim_weight[0]
276
- output_size = np.append(input_size[0:-1], dim_weight[0])
277
- output_name = node.output[0]
278
- return macs, output_size, output_name
279
- pass
280
-
281
-
282
- def onnx_counter_maxpool(diction, node):
283
- """Calculate MACs and output size for ONNX MaxPool operation based on input node attributes and dimensions."""
284
- # print(node)
285
- macs = calculate_zero_ops()
286
- output_name = node.output[0]
287
- dim_pad = None
288
- for attr in node.attribute:
289
- # print(attr)
290
- if attr.name == "kernel_shape":
291
- dim_kernel = attr.ints # kw,kh
292
- elif attr.name == "strides":
293
- dim_stride = attr.ints
294
- elif attr.name == "pads":
295
- dim_pad = attr.ints
296
- elif attr.name == "dilations":
297
- dim_dil = attr.ints
298
- # print(dim_dil)
299
- dim_input = diction[node.input[0]]
300
- hw = dim_input[-np.array(dim_kernel).size :]
301
- if dim_pad is not None:
302
- for i in range(hw.size):
303
- hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_kernel[i]) / dim_stride[i] + 1)
304
- output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw)
305
- else:
306
- for i in range(hw.size):
307
- hw[i] = int((hw[i] - dim_kernel[i]) / dim_stride[i] + 1)
308
- output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw)
309
- # print(macs, output_size, output_name)
310
- return macs, output_size, output_name
311
-
312
-
313
- def onnx_counter_globalaveragepool(diction, node):
314
- """Counts MACs and computes output size for a global average pooling layer in an ONNX model."""
315
- macs = calculate_zero_ops()
316
- output_name = node.output[0]
317
- input_size = diction[node.input[0]]
318
- output_size = input_size
319
- return macs, output_size, output_name
320
-
321
-
322
- def onnx_counter_concat(diction, node):
323
- """Counts MACs and computes output size for a concatenation layer along a specified axis in an ONNX model."""
324
- # print(diction[node.input[0]])
325
- axis = node.attribute[0].i
326
- input_size = diction[node.input[0]]
327
- for i in node.input:
328
- dim_concat = diction[i][axis]
329
- output_size = input_size
330
- output_size[axis] = dim_concat
331
- output_name = node.output[0]
332
- macs = calculate_zero_ops()
333
- return macs, output_size, output_name
334
-
335
-
336
- def onnx_counter_clip(diction, node):
337
- """Calculate MACs, output size, and output name for an ONNX node clip operation using provided dimensions and input
338
- size.
339
- """
340
- macs = calculate_zero_ops()
341
- output_name = node.output[0]
342
- input_size = diction[node.input[0]]
343
- output_size = input_size
344
- return macs, output_size, output_name
345
-
346
-
347
- onnx_operators = {
348
- "MatMul": onnx_counter_matmul,
349
- "Add": onnx_counter_add,
350
- "Conv": onnx_counter_conv,
351
- "Mul": onnx_counter_mul,
352
- "Constant": onnx_counter_constant,
353
- "BatchNormalization": onnx_counter_bn,
354
- "Relu": onnx_counter_relu,
355
- "ReduceMean": onnx_counter_reducemean,
356
- "Sub": onnx_counter_sub,
357
- "Pow": onnx_counter_pow,
358
- "Sqrt": onnx_counter_sqrt,
359
- "Div": onnx_counter_div,
360
- "InstanceNormalization": onnx_counter_instance,
361
- "Softmax": onnx_counter_softmax,
362
- "Pad": onnx_counter_pad,
363
- "AveragePool": onnx_counter_averagepool,
364
- "MaxPool": onnx_counter_maxpool,
365
- "Flatten": onnx_counter_flatten,
366
- "Gemm": onnx_counter_gemm,
367
- "GlobalAveragePool": onnx_counter_globalaveragepool,
368
- "Concat": onnx_counter_concat,
369
- "Clip": onnx_counter_clip,
370
- None: None,
371
- }