From 296d04b58524e16961652cf90205f95cfa408d38 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 1 Jun 2024 03:42:08 +0200 Subject: [PATCH] Refactor code (#18) Co-authored-by: UltralyticsAssistant --- benchmark/README.md | 9 +- tests/test_conv2d.py | 8 +- tests/test_matmul.py | 1 - tests/test_relu.py | 1 - tests/test_utils.py | 2 - thop/__init__.py | 2 +- thop/fx_profile.py | 8 +- thop/onnx_profile.py | 82 -------- thop/profile.py | 10 +- thop/rnn_hooks.py | 69 ++++--- thop/utils.py | 13 +- thop/vision/basic_hooks.py | 4 +- thop/vision/calc_func.py | 22 +-- thop/vision/efficientnet.py | 9 - thop/vision/onnx_counter.py | 371 ------------------------------------ 15 files changed, 58 insertions(+), 553 deletions(-) delete mode 100644 thop/onnx_profile.py delete mode 100644 thop/vision/efficientnet.py delete mode 100644 thop/vision/onnx_counter.py diff --git a/benchmark/README.md b/benchmark/README.md index 89b834c..20756bd 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -11,16 +11,17 @@ However, the application in real world is far more complex. Let's consider a mat ```python for i in range(m): for j in range(n): - C[i][j] += A[i][j] * B[j] # one mul-add + C[i][j] += A[i][j] * B[j] # one mul-add ``` It would be `mn` `MACs` and `2mn` `FLOPs`. But such implementation is slow and parallelization is necessary to run faster ```python for i in range(m): - parallelfor j in range(n): - d[j] = A[i][j] * B[j] # one mul - C[i][j] = sum(d) # n adds + parallelfor + j in range(n): + d[j] = A[i][j] * B[j] # one mul +C[i][j] = sum(d) # n adds ``` Then the number of `MACs` is no longer `mn` . diff --git a/tests/test_conv2d.py b/tests/test_conv2d.py index 75d0c96..e058e25 100644 --- a/tests/test_conv2d.py +++ b/tests/test_conv2d.py @@ -1,7 +1,5 @@ -import pytest import torch import torch.nn as nn -from jinja2 import StrictUndefined from thop import profile @@ -22,7 +20,7 @@ def test_conv2d_no_bias(self): _, _, oh, ow = out.shape flops, params = profile(net, inputs=(data,)) - assert flops == 810000, f"{flops} v.s. {810000}" + assert flops == 810000, f"{flops} v.s. 810000" def test_conv2d(self): """Tests a Conv2D layer with specific input dimensions, kernel size, stride, padding, dilation, and groups.""" @@ -37,11 +35,11 @@ def test_conv2d(self): _, _, oh, ow = out.shape flops, params = profile(net, inputs=(data,)) - assert flops == 810000, f"{flops} v.s. {810000}" + assert flops == 810000, f"{flops} v.s. 810000" def test_conv2d_random(self): """Test Conv2D layer with random parameters and validate the computed FLOPs and parameters using 'profile'.""" - for i in range(10): + for _ in range(10): out_c, kh, kw = torch.randint(1, 20, (3,)).tolist() n, in_c, ih, iw = torch.randint(1, 20, (4,)).tolist() # torch.randint(1, 10, (4,)).tolist() ih += kh diff --git a/tests/test_matmul.py b/tests/test_matmul.py index 15fe893..48391c9 100644 --- a/tests/test_matmul.py +++ b/tests/test_matmul.py @@ -1,4 +1,3 @@ -import pytest import torch import torch.nn as nn diff --git a/tests/test_relu.py b/tests/test_relu.py index 1fcffd0..f131bf5 100644 --- a/tests/test_relu.py +++ b/tests/test_relu.py @@ -1,4 +1,3 @@ -import pytest import torch import torch.nn as nn diff --git a/tests/test_utils.py b/tests/test_utils.py index 53b69dd..0358e7e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,5 +1,3 @@ -import pytest - from thop import utils diff --git a/thop/__init__.py b/thop/__init__.py index 3f9851f..6e2189d 100644 --- a/thop/__init__.py +++ b/thop/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.2.4" +__version__ = "0.2.5" import torch diff --git a/thop/fx_profile.py b/thop/fx_profile.py index ea15050..52a6c29 100644 --- a/thop/fx_profile.py +++ b/thop/fx_profile.py @@ -85,9 +85,7 @@ def count_nn_bn2d(module: nn.BatchNorm2d, input_shapes, output_shapes): """Calculate the total operations for a given nn.BatchNorm2d module based on its output shape.""" assert len(output_shapes) == 1, "nn.BatchNorm2d should only have one output" y = output_shapes[0] - # y = (x - mean) / \sqrt{var + e} * weight + bias - total_ops = 2 * y.numel() - return total_ops + return 2 * y.numel() zero_ops = ( @@ -120,7 +118,7 @@ def count_nn_bn2d(module: nn.BatchNorm2d, input_shapes, output_shapes): from torch.fx import symbolic_trace from torch.fx.passes.shape_prop import ShapeProp -from .utils import prGreen, prRed, prYellow +from .utils import prRed, prYellow def null_print(*args, **kwargs): @@ -193,7 +191,7 @@ def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False): prRed(f"{key} is missing") print("module type:", type(m)) if isinstance(m, zero_ops): - print(f"weight_shape: None") + print("weight_shape: None") else: print(type(m)) print(f"weight_shape: {mod.state_dict()[node.target + '.weight'].shape}") diff --git a/thop/onnx_profile.py b/thop/onnx_profile.py deleted file mode 100644 index 926d1d5..0000000 --- a/thop/onnx_profile.py +++ /dev/null @@ -1,82 +0,0 @@ -import numpy as np -import onnx -import torch -import torch.nn -from onnx import numpy_helper - -from thop.vision.onnx_counter import onnx_operators - - -class OnnxProfile: - def __init__(self): - """Initialize the OnnxProfile class with necessary imports for ONNX profiling.""" - pass - - def calculate_params(self, model: onnx.ModelProto): - """Calculate the total number of parameters in an ONNX model.""" - onnx_weights = model.graph.initializer - params = 0 - - for onnx_w in onnx_weights: - try: - weight = numpy_helper.to_array(onnx_w) - params += np.prod(weight.shape) - except Exception as _: - pass - - return params - - def create_dict(self, weight, input, output): - """Create and return a dictionary mapping weight, input, and output names to their respective dimensions.""" - diction = {} - for w in weight: - dim = np.array(w.dims) - diction[str(w.name)] = dim - if dim.size == 1: - diction[str(w.name)] = np.append(1, dim) - for i in input: - # print(i.type.tensor_type.shape.dim[0].dim_value) - dim = np.array(i.type.tensor_type.shape.dim[0].dim_value) - # print(i.type.tensor_type.shape.dim.__sizeof__()) - # name2dims[str(i.name)] = [dim] - dim = [] - for key in i.type.tensor_type.shape.dim: - dim = np.append(dim, int(key.dim_value)) - # print(key.dim_value) - # print(dim) - diction[str(i.name)] = dim - if dim.size == 1: - diction[str(i.name)] = np.append(1, dim) - for o in output: - dim = np.array(o.type.tensor_type.shape.dim[0].dim_value) - diction[str(o.name)] = [dim] - if dim.size == 1: - diction[str(o.name)] = np.append(1, dim) - return diction - - def nodes_counter(self, diction, node): - """Count nodes of a specific type in an ONNX graph, returning the count and associated node operation - details. - """ - if node.op_type not in onnx_operators: - print("Sorry, we haven't add ", node.op_type, "into dictionary.") - return 0, None, None - else: - fn = onnx_operators[node.op_type] - return fn(diction, node) - - def calculate_macs(self, model: onnx.ModelProto) -> torch.DoubleTensor: - macs = 0 - name2dims = {} - weight = model.graph.initializer - nodes = model.graph.node - input = model.graph.input - output = model.graph.output - name2dims = self.create_dict(weight, input, output) - macs = 0 - for n in nodes: - macs_adding, out_size, outname = self.nodes_counter(name2dims, n) - - name2dims[outname] = out_size - macs += macs_adding - return np.array(macs[0]) diff --git a/thop/profile.py b/thop/profile.py index 9f2c92d..c8846b3 100644 --- a/thop/profile.py +++ b/thop/profile.py @@ -1,9 +1,7 @@ from thop.rnn_hooks import * from thop.vision.basic_hooks import * -# logger = logging.getLogger(__name__) -# logger.setLevel(logging.INFO) -from .utils import prGreen, prRed, prYellow +from .utils import prRed default_dtype = torch.float64 @@ -68,7 +66,7 @@ def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing= verbose = True def add_hooks(m): - if len(list(m.children())) > 0: + if list(m.children()): return if hasattr(m, "total_ops") or hasattr(m, "total_params"): @@ -114,7 +112,7 @@ def add_hooks(m): total_ops = 0 total_params = 0 for m in model.modules(): - if len(list(m.children())) > 0: # skip for non-leaf module + if list(m.children()): # skip for non-leaf module continue total_ops += m.total_ops total_params += m.total_params @@ -129,7 +127,7 @@ def add_hooks(m): # remove temporal buffers for n, m in model.named_modules(): - if len(list(m.children())) > 0: + if list(m.children()): continue if "total_ops" in m._buffers: m._buffers.pop("total_ops") diff --git a/thop/rnn_hooks.py b/thop/rnn_hooks.py index 471aed5..d87186f 100644 --- a/thop/rnn_hooks.py +++ b/thop/rnn_hooks.py @@ -103,13 +103,12 @@ def count_rnn(m: nn.RNN, x, y): if isinstance(x[0], PackedSequence): batch_size = torch.max(x[0].batch_sizes) num_steps = x[0].batch_sizes.size(0) + elif m.batch_first: + batch_size = x[0].size(0) + num_steps = x[0].size(1) else: - if m.batch_first: - batch_size = x[0].size(0) - num_steps = x[0].size(1) - else: - batch_size = x[0].size(1) - num_steps = x[0].size(0) + batch_size = x[0].size(1) + num_steps = x[0].size(0) total_ops = 0 if m.bidirectional: @@ -117,12 +116,12 @@ def count_rnn(m: nn.RNN, x, y): else: total_ops += _count_rnn_cell(input_size, hidden_size, bias) - for i in range(num_layers - 1): - if m.bidirectional: - total_ops += _count_rnn_cell(hidden_size * 2, hidden_size, bias) * 2 - else: - total_ops += _count_rnn_cell(hidden_size, hidden_size, bias) - + for _ in range(num_layers - 1): + total_ops += ( + _count_rnn_cell(hidden_size * 2, hidden_size, bias) * 2 + if m.bidirectional + else _count_rnn_cell(hidden_size, hidden_size, bias) + ) # time unroll total_ops *= num_steps # batch_size @@ -141,13 +140,12 @@ def count_gru(m: nn.GRU, x, y): if isinstance(x[0], PackedSequence): batch_size = torch.max(x[0].batch_sizes) num_steps = x[0].batch_sizes.size(0) + elif m.batch_first: + batch_size = x[0].size(0) + num_steps = x[0].size(1) else: - if m.batch_first: - batch_size = x[0].size(0) - num_steps = x[0].size(1) - else: - batch_size = x[0].size(1) - num_steps = x[0].size(0) + batch_size = x[0].size(1) + num_steps = x[0].size(0) total_ops = 0 if m.bidirectional: @@ -155,12 +153,12 @@ def count_gru(m: nn.GRU, x, y): else: total_ops += _count_gru_cell(input_size, hidden_size, bias) - for i in range(num_layers - 1): - if m.bidirectional: - total_ops += _count_gru_cell(hidden_size * 2, hidden_size, bias) * 2 - else: - total_ops += _count_gru_cell(hidden_size, hidden_size, bias) - + for _ in range(num_layers - 1): + total_ops += ( + _count_gru_cell(hidden_size * 2, hidden_size, bias) * 2 + if m.bidirectional + else _count_gru_cell(hidden_size, hidden_size, bias) + ) # time unroll total_ops *= num_steps # batch_size @@ -181,13 +179,12 @@ def count_lstm(m: nn.LSTM, x, y): if isinstance(x[0], PackedSequence): batch_size = torch.max(x[0].batch_sizes) num_steps = x[0].batch_sizes.size(0) + elif m.batch_first: + batch_size = x[0].size(0) + num_steps = x[0].size(1) else: - if m.batch_first: - batch_size = x[0].size(0) - num_steps = x[0].size(1) - else: - batch_size = x[0].size(1) - num_steps = x[0].size(0) + batch_size = x[0].size(1) + num_steps = x[0].size(0) total_ops = 0 if m.bidirectional: @@ -195,12 +192,12 @@ def count_lstm(m: nn.LSTM, x, y): else: total_ops += _count_lstm_cell(input_size, hidden_size, bias) - for i in range(num_layers - 1): - if m.bidirectional: - total_ops += _count_lstm_cell(hidden_size * 2, hidden_size, bias) * 2 - else: - total_ops += _count_lstm_cell(hidden_size, hidden_size, bias) - + for _ in range(num_layers - 1): + total_ops += ( + _count_lstm_cell(hidden_size * 2, hidden_size, bias) * 2 + if m.bidirectional + else _count_lstm_cell(hidden_size, hidden_size, bias) + ) # time unroll total_ops *= num_steps # batch_size diff --git a/thop/utils.py b/thop/utils.py index 8f324e4..04114d1 100644 --- a/thop/utils.py +++ b/thop/utils.py @@ -20,15 +20,6 @@ def actual_call(*args, **kwargs): prGreen = colorful_print(print, color=COLOR_GREEN) prYellow = colorful_print(print, color=COLOR_YELLOW) -# def prRed(skk): -# print("\033[91m{}\033[00m".format(skk)) - -# def prGreen(skk): -# print("\033[92m{}\033[00m".format(skk)) - -# def prYellow(skk): -# print("\033[93m{}\033[00m".format(skk)) - def clever_format(nums, format="%.2f"): """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"): else: clever_nums.append(format % num + "B") - clever_nums = clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,) - - return clever_nums + return clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,) if __name__ == "__main__": diff --git a/thop/vision/basic_hooks.py b/thop/vision/basic_hooks.py index cc86018..5a7897a 100644 --- a/thop/vision/basic_hooks.py +++ b/thop/vision/basic_hooks.py @@ -1,7 +1,5 @@ -import argparse import logging -import torch import torch.nn as nn from torch.nn.modules.conv import _ConvNd @@ -139,7 +137,7 @@ def count_upsample(m, x, y): "bilinear", "bicubic", ): # "trilinear" - logging.warning("mode %s is not implemented yet, take it a zero op" % m.mode) + logging.warning(f"mode {m.mode} is not implemented yet, take it a zero op") m.total_ops += 0 else: x = x[0] diff --git a/thop/vision/calc_func.py b/thop/vision/calc_func.py index ac1e7b3..9b9f3d0 100644 --- a/thop/vision/calc_func.py +++ b/thop/vision/calc_func.py @@ -14,23 +14,17 @@ def l_prod(in_list): def l_sum(in_list): """Calculate the sum of all elements in a list.""" - res = 0 - for _ in in_list: - res += _ - return res + return sum(in_list) def calculate_parameters(param_list): """Calculate the total number of parameters in a list of tensors.""" - total_params = 0 - for p in param_list: - total_params += torch.DoubleTensor([p.nelement()]) - return total_params + return sum(torch.DoubleTensor([p.nelement()]) for p in param_list) def calculate_zero_ops(): """Return a tensor initialized to zero.""" - return torch.DoubleTensor([int(0)]) + return torch.DoubleTensor([0]) 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): def calculate_upsample(mode: str, output_size): """Calculate the number of operations for upsample methods given the mode and output size.""" total_ops = output_size - if mode == "linear": - total_ops *= 5 + if mode == "bicubic": + total_ops *= 224 + 35 elif mode == "bilinear": total_ops *= 11 - elif mode == "bicubic": - ops_solve_A = 224 # 128 muls + 96 adds - ops_solve_p = 35 # 16 muls + 12 adds + 4 muls + 3 adds - total_ops *= ops_solve_A + ops_solve_p + elif mode == "linear": + total_ops *= 5 elif mode == "trilinear": total_ops *= 13 * 2 + 5 return torch.DoubleTensor([int(total_ops)]) diff --git a/thop/vision/efficientnet.py b/thop/vision/efficientnet.py deleted file mode 100644 index 8de88d1..0000000 --- a/thop/vision/efficientnet.py +++ /dev/null @@ -1,9 +0,0 @@ -import argparse -import logging - -import torch -import torch.nn as nn -from efficientnet_pytorch.utils import Conv2dDynamicSamePadding, Conv2dStaticSamePadding -from torch.nn.modules.conv import _ConvNd - -register_hooks = {} diff --git a/thop/vision/onnx_counter.py b/thop/vision/onnx_counter.py deleted file mode 100644 index 4189013..0000000 --- a/thop/vision/onnx_counter.py +++ /dev/null @@ -1,371 +0,0 @@ -import numpy as np -import torch -from onnx import numpy_helper - -from thop.vision.basic_hooks import zero_ops - -from .calc_func import ( - calculate_avgpool, - calculate_conv, - calculate_norm, - calculate_softmax, - calculate_zero_ops, - counter_div, - counter_matmul, - counter_mul, - counter_pow, - counter_sqrt, -) - - -def onnx_counter_matmul(diction, node): - """Calculates multiply-accumulate operations and output size for matrix multiplication in an ONNX model node.""" - input1 = node.input[0] - input2 = node.input[1] - input1_dim = diction[input1] - input2_dim = diction[input2] - out_size = np.append(input1_dim[0:-1], input2_dim[-1]) - output_name = node.output[0] - macs = counter_matmul(input1_dim, out_size[-2:]) - return macs, out_size, output_name - - -def onnx_counter_add(diction, node): - """Calculate multiply-accumulate operations (MACs), output size, and output name for ONNX addition nodes.""" - if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size: - out_size = diction[node.input[1]] - else: - out_size = diction[node.input[0]] - output_name = node.output[0] - macs = calculate_zero_ops() - # if '140' in diction: - # print(diction['140'],output_name) - return macs, out_size, output_name - - -def onnx_counter_conv(diction, node): - """Calculates MACs, output size, and name for an ONNX convolution node based on input tensor dimensions and node - attributes. - """ - # bias,kernelsize,outputsize - dim_bias = 0 - input_count = 0 - for i in node.input: - input_count += 1 - if input_count == 3: - dim_bias = 1 - dim_weight = diction[node.input[1]] - else: - dim_weight = diction[node.input[1]] - for attr in node.attribute: - # print(attr) - if attr.name == "kernel_shape": - dim_kernel = attr.ints # kw,kh - if attr.name == "strides": - dim_stride = attr.ints - if attr.name == "pads": - dim_pad = attr.ints - if attr.name == "dilations": - dim_dil = attr.ints - if attr.name == "group": - group = attr.i - # print(dim_dil) - dim_input = diction[node.input[0]] - output_size = np.append(dim_input[0 : -np.array(dim_kernel).size - 1], dim_weight[0]) - hw = np.array(dim_input[-np.array(dim_kernel).size :]) - for i in range(hw.size): - hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_dil[i] * (dim_kernel[i] - 1) - 1) / dim_stride[i] + 1) - output_size = np.append(output_size, hw) - macs = calculate_conv(dim_bias, np.prod(dim_kernel), np.prod(output_size), dim_weight[1], group) - output_name = node.output[0] - - # if '140' in diction: - # print("conv",diction['140'],output_name) - return macs, output_size, output_name - - -def onnx_counter_constant(diction, node): - """Calculate MACs, output size, and output name for a constant operation in an ONNX model.""" - macs = calculate_zero_ops() - output_name = node.output[0] - output_size = [1] - # print(macs, output_size, output_name) - return macs, output_size, output_name - - -def onnx_counter_mul(diction, node): - """Calculate MACs, output size, and output name for a multiplication operation in an ONNX model.""" - if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size: - input_size = diction[node.input[1]] - else: - input_size = diction[node.input[0]] - macs = counter_mul(np.prod(input_size)) - output_size = diction[node.input[0]] - output_name = node.output[0] - return macs, output_size, output_name - - -def onnx_counter_bn(diction, node): - """Calculates MACs, output size, and output name for batch normalization layers in an ONNX model.""" - input_size = diction[node.input[0]] - macs = calculate_norm(np.prod(input_size)) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_relu(diction, node): - """Calculates MACs, output size, and output name for ReLU layers in an ONNX model.""" - input_size = diction[node.input[0]] - macs = calculate_zero_ops() - output_name = node.output[0] - output_size = input_size - # print(macs, output_size, output_name) - # if '140' in diction: - # print("relu",diction['140'],output_name) - return macs, output_size, output_name - - -def onnx_counter_reducemean(diction, node): - """Compute MACs, output size, and name for the ReduceMean ONNX node, adjusting dimensions based on the 'axes' and - 'keepdims' attributes. - """ - keep_dim = 0 - for attr in node.attribute: - if "axes" in attr.name: - dim_axis = np.array(attr.ints) - elif "keepdims" in attr.name: - keep_dim = attr.i - - input_size = diction[node.input[0]] - macs = calculate_zero_ops() - output_name = node.output[0] - if keep_dim == 1: - output_size = input_size - else: - output_size = np.delete(input_size, dim_axis) - # output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_sub(diction, node): - """Computes MACs, output size, and output name for a given ONNX node with specified input size.""" - input_size = diction[node.input[0]] - macs = calculate_zero_ops() - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_pow(diction, node): - """Calculates MACs, output size, and output name for a given ONNX 'Pow' node with specified input size.""" - if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size: - input_size = diction[node.input[1]] - else: - input_size = diction[node.input[0]] - macs = counter_pow(np.prod(input_size)) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_sqrt(diction, node): - """Calculate MACs and output information for the SQRT operation in an ONNX node.""" - input_size = diction[node.input[0]] - macs = counter_sqrt(np.prod(input_size)) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_div(diction, node): - """Calculate MACs and output information for the DIV operation in an ONNX node.""" - if np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size: - input_size = diction[node.input[1]] - else: - input_size = diction[node.input[0]] - macs = counter_div(np.prod(input_size)) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_instance(diction, node): - """Calculate MACs, output size, and name for an ONNX node instance.""" - input_size = diction[node.input[0]] - macs = calculate_norm(np.prod(input_size)) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_softmax(diction, node): - """Calculate MACs, output size, and name for an ONNX softmax node instance.""" - input_size = diction[node.input[0]] - dim = node.attribute[0].i - nfeatures = input_size[dim] - batch_size = np.prod(input_size) / nfeatures - macs = calculate_softmax(nfeatures, batch_size) - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_pad(diction, node): - """Compute memory access cost (MACs), output size, and output name for ONNX pad operation.""" - # if - # if (np.array(diction[node.input[1]]).size >= np.array(diction[node.input[0]]).size): - # input_size = diction[node.input[1]] - # else: - # input_size = diction[node.input[0]] - input_size = diction[node.input[0]] - macs = calculate_zero_ops() - output_name = node.output[0] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_averagepool(diction, node): - """Calculate MACs and output size for an AveragePool ONNX operation based on input dimensions and attributes.""" - macs = calculate_avgpool(np.prod(diction[node.input[0]])) - output_name = node.output[0] - dim_pad = None - for attr in node.attribute: - # print(attr) - if attr.name == "kernel_shape": - dim_kernel = attr.ints # kw,kh - elif attr.name == "strides": - dim_stride = attr.ints - elif attr.name == "pads": - dim_pad = attr.ints - elif attr.name == "dilations": - dim_dil = attr.ints - # print(dim_dil) - dim_input = diction[node.input[0]] - hw = dim_input[-np.array(dim_kernel).size :] - if dim_pad is not None: - for i in range(hw.size): - hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_kernel[i]) / dim_stride[i] + 1) - output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw) - else: - for i in range(hw.size): - hw[i] = int((hw[i] - dim_kernel[i]) / dim_stride[i] + 1) - output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw) - # print(macs, output_size, output_name) - return macs, output_size, output_name - - -def onnx_counter_flatten(diction, node): - """Returns MACs, output size, and output name for an ONNX Flatten node.""" - macs = calculate_zero_ops() - output_name = node.output[0] - axis = node.attribute[0].i - input_size = diction[node.input[0]] - output_size = np.append(input_size[axis - 1], np.prod(input_size[axis:])) - # print("flatten",output_size) - return macs, output_size, output_name - - -def onnx_counter_gemm(diction, node): - """Calculate multiply–accumulate operations (MACs), output size, and name for ONNX Gemm node.""" - # Compute Y = alpha * A' * B' + beta * C - input_size = diction[node.input[0]] - dim_weight = diction[node.input[1]] - # print(input_size,dim_weight) - macs = np.prod(input_size) * dim_weight[1] + dim_weight[0] - output_size = np.append(input_size[0:-1], dim_weight[0]) - output_name = node.output[0] - return macs, output_size, output_name - pass - - -def onnx_counter_maxpool(diction, node): - """Calculate MACs and output size for ONNX MaxPool operation based on input node attributes and dimensions.""" - # print(node) - macs = calculate_zero_ops() - output_name = node.output[0] - dim_pad = None - for attr in node.attribute: - # print(attr) - if attr.name == "kernel_shape": - dim_kernel = attr.ints # kw,kh - elif attr.name == "strides": - dim_stride = attr.ints - elif attr.name == "pads": - dim_pad = attr.ints - elif attr.name == "dilations": - dim_dil = attr.ints - # print(dim_dil) - dim_input = diction[node.input[0]] - hw = dim_input[-np.array(dim_kernel).size :] - if dim_pad is not None: - for i in range(hw.size): - hw[i] = int((hw[i] + 2 * dim_pad[i] - dim_kernel[i]) / dim_stride[i] + 1) - output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw) - else: - for i in range(hw.size): - hw[i] = int((hw[i] - dim_kernel[i]) / dim_stride[i] + 1) - output_size = np.append(dim_input[0 : -np.array(dim_kernel).size], hw) - # print(macs, output_size, output_name) - return macs, output_size, output_name - - -def onnx_counter_globalaveragepool(diction, node): - """Counts MACs and computes output size for a global average pooling layer in an ONNX model.""" - macs = calculate_zero_ops() - output_name = node.output[0] - input_size = diction[node.input[0]] - output_size = input_size - return macs, output_size, output_name - - -def onnx_counter_concat(diction, node): - """Counts MACs and computes output size for a concatenation layer along a specified axis in an ONNX model.""" - # print(diction[node.input[0]]) - axis = node.attribute[0].i - input_size = diction[node.input[0]] - for i in node.input: - dim_concat = diction[i][axis] - output_size = input_size - output_size[axis] = dim_concat - output_name = node.output[0] - macs = calculate_zero_ops() - return macs, output_size, output_name - - -def onnx_counter_clip(diction, node): - """Calculate MACs, output size, and output name for an ONNX node clip operation using provided dimensions and input - size. - """ - macs = calculate_zero_ops() - output_name = node.output[0] - input_size = diction[node.input[0]] - output_size = input_size - return macs, output_size, output_name - - -onnx_operators = { - "MatMul": onnx_counter_matmul, - "Add": onnx_counter_add, - "Conv": onnx_counter_conv, - "Mul": onnx_counter_mul, - "Constant": onnx_counter_constant, - "BatchNormalization": onnx_counter_bn, - "Relu": onnx_counter_relu, - "ReduceMean": onnx_counter_reducemean, - "Sub": onnx_counter_sub, - "Pow": onnx_counter_pow, - "Sqrt": onnx_counter_sqrt, - "Div": onnx_counter_div, - "InstanceNormalization": onnx_counter_instance, - "Softmax": onnx_counter_softmax, - "Pad": onnx_counter_pad, - "AveragePool": onnx_counter_averagepool, - "MaxPool": onnx_counter_maxpool, - "Flatten": onnx_counter_flatten, - "Gemm": onnx_counter_gemm, - "GlobalAveragePool": onnx_counter_globalaveragepool, - "Concat": onnx_counter_concat, - "Clip": onnx_counter_clip, - None: None, -}