ultralytics-thop 0.0.1__tar.gz → 0.0.2__tar.gz

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics-thop
-Version: 0.0.1
+Version: 0.0.2
 Summary: A tool to count the FLOPs of PyTorch model.
 Author-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>
 Maintainer-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ultralytics-thop"
-version = "0.0.1"  # Placeholder version, needs to be dynamically set
+version = "0.0.2"  # Placeholder version, needs to be dynamically set
 description = "A tool to count the FLOPs of PyTorch model."
 readme = "README.md"
 requires-python = ">=3.8"

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/tests/test_conv2d.py

@@ -8,6 +8,9 @@ from thop import profile
 
 class TestUtils:
     def test_conv2d_no_bias(self):
+        """Tests a 2D convolutional layer without bias using THOP profiling with predefined input dimensions and
+        convolution parameters.
+        """
         n, in_c, ih, iw = 1, 3, 32, 32  # torch.randint(1, 10, (4,)).tolist()
         out_c, kh, kw = 12, 5, 5
         s, p, d, g = 1, 1, 1, 1
@@ -22,6 +25,7 @@ class TestUtils:
         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."""
         n, in_c, ih, iw = 1, 3, 32, 32  # torch.randint(1, 10, (4,)).tolist()
         out_c, kh, kw = 12, 5, 5
         s, p, d, g = 1, 1, 1, 1
@@ -36,6 +40,7 @@ class TestUtils:
         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):
             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()

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/tests/test_matmul.py

@@ -7,6 +7,7 @@ from thop import profile
 
 class TestUtils:
     def test_matmul_case2(self):
+        """Test matrix multiplication case asserting the FLOPs and parameters of a nn.Linear layer."""
         n, in_c, out_c = 1, 100, 200
         net = nn.Linear(in_c, out_c)
         flops, params = profile(net, inputs=(torch.randn(n, in_c),))
@@ -14,6 +15,7 @@ class TestUtils:
         assert flops == n * in_c * out_c
 
     def test_matmul_case2(self):
+        """Tests matrix multiplication to assert FLOPs and parameters of nn.Linear layer using random dimensions."""
         for i in range(10):
             n, in_c, out_c = torch.randint(1, 500, (3,)).tolist()
             net = nn.Linear(in_c, out_c)
@@ -22,6 +24,7 @@ class TestUtils:
             assert flops == n * in_c * out_c
 
     def test_conv2d(self):
+        """Tests the number of FLOPs and parameters for a randomly initialized nn.Linear layer using torch.profiler."""
         n, in_c, out_c = torch.randint(1, 500, (3,)).tolist()
         net = nn.Linear(in_c, out_c)
         flops, params = profile(net, inputs=(torch.randn(n, in_c),))

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/tests/test_relu.py

@@ -7,6 +7,9 @@ from thop import profile
 
 class TestUtils:
     def test_relu(self):
+        """Tests the ReLU activation function to ensure it has zero FLOPs and checks parameter count using THOP
+        profiling.
+        """
         n, in_c, out_c = 1, 100, 200
         data = torch.randn(n, in_c)
         net = nn.ReLU()

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/tests/test_utils.py

@@ -5,12 +5,16 @@ from thop import utils
 
 class TestUtils:
     def test_clever_format_returns_formatted_number(self):
+        """Tests that the clever_format function returns a formatted number string with a '1.00B' pattern."""
         nums = 1
         format = "%.2f"
         clever_nums = utils.clever_format(nums, format)
         assert clever_nums == "1.00B"
 
     def test_clever_format_returns_formatted_numbers(self):
+        """Tests that the clever_format function correctly formats a list of numbers as strings with a '1.00B'
+        pattern.
+        """
         nums = [1, 2]
         format = "%.2f"
         clever_nums = utils.clever_format(nums, format)

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/fx_profile.py

@@ -13,15 +13,17 @@ if LooseVersion(torch.__version__) < LooseVersion("1.8.0"):
 
 
 def count_clamp(input_shapes, output_shapes):
+    """Ensures proper array sizes for tensors by clamping input and output shapes."""
     return 0
 
 
 def count_mul(input_shapes, output_shapes):
-    # element-wise
+    """Returns the number of elements in the first output shape."""
     return output_shapes[0].numel()
 
 
 def count_matmul(input_shapes, output_shapes):
+    """Calculates the total number of operations for a matrix multiplication given input and output shapes."""
     in_shape = input_shapes[0]
     out_shape = output_shapes[0]
     in_features = in_shape[-1]
@@ -30,6 +32,7 @@ def count_matmul(input_shapes, output_shapes):
 
 
 def count_fn_linear(input_shapes, output_shapes, *args, **kwargs):
+    """Calculates total operations (FLOPs) for a linear layer given input and output shapes."""
     mul_flops = count_matmul(input_shapes, output_shapes)
     if "bias" in kwargs:
         add_flops = output_shapes[0].numel()
@@ -40,6 +43,7 @@ from .vision.calc_func import calculate_conv
 
 
 def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs):
+    """Calculates total operations (FLOPs) for a 2D convolutional layer given input and output shapes."""
     inputs, weight, bias, stride, padding, dilation, groups = args
     if len(input_shapes) == 2:
         x_shape, k_shape = input_shapes
@@ -56,14 +60,17 @@ def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs):
 
 
 def count_nn_linear(module: nn.Module, input_shapes, output_shapes):
+    """Counts the FLOPs for a fully connected (linear) layer in a neural network module."""
     return count_matmul(input_shapes, output_shapes)
 
 
 def count_zero_ops(module: nn.Module, input_shapes, output_shapes, *args, **kwargs):
+    """Returns 0 for the given neural network module, input shapes, and output shapes."""
     return 0
 
 
 def count_nn_conv2d(module: nn.Conv2d, input_shapes, output_shapes):
+    """Calculates total operations for a 2D convolutional neural network layer in a given neural network module."""
     bias_op = 1 if module.bias is not None else 0
     out_shape = output_shapes[0]
 
@@ -75,6 +82,7 @@ def count_nn_conv2d(module: nn.Conv2d, input_shapes, output_shapes):
 
 
 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
@@ -116,10 +124,14 @@ from .utils import prGreen, prRed, prYellow
 
 
 def null_print(*args, **kwargs):
+    """A no-op print function that takes any arguments without performing any actions."""
     return
 
 
 def fx_profile(mod: nn.Module, input: th.Tensor, verbose=False):
+    """Profiles the given torch.nn Module to calculate total FLOPs for each operation and prints detailed node
+    information if verbose.
+    """
     gm: torch.fx.GraphModule = symbolic_trace(mod)
     g = gm.graph
     ShapeProp(gm).propagate(input)
@@ -204,16 +216,21 @@ if __name__ == "__main__":
 
     class MyOP(nn.Module):
         def forward(self, input):
+            """Performs forward pass on given input data."""
             return input / 1
 
     class MyModule(torch.nn.Module):
         def __init__(self):
+            """Initializes MyModule with two linear layers and a custom MyOP operator."""
             super().__init__()
             self.linear1 = torch.nn.Linear(5, 3)
             self.linear2 = torch.nn.Linear(5, 3)
             self.myop = MyOP()
 
         def forward(self, x):
+            """Applies two linear transformations to the input tensor, clamps the second, then combines and processes
+            with MyOP operator.
+            """
             out1 = self.linear1(x)
             out2 = self.linear2(x).clamp(min=0.0, max=1.0)
             return self.myop(out1 + out2)

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/onnx_profile.py

@@ -9,9 +9,11 @@ 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
 
@@ -25,6 +27,7 @@ class OnnxProfile:
         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)
@@ -52,6 +55,9 @@ class OnnxProfile:
         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

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/profile.py

@@ -66,6 +66,9 @@ if LooseVersion(torch.__version__) >= LooseVersion("1.1.0"):
 
 
 def profile_origin(model, inputs, custom_ops=None, verbose=True, report_missing=False):
+    """Profiles a PyTorch model's operations and parameters by applying custom or default hooks and returns total
+    operations and parameters.
+    """
     handler_collection = []
     types_collection = set()
     if custom_ops is None:
@@ -162,6 +165,7 @@ def profile(
         verbose = True
 
     def add_hooks(m: nn.Module):
+        """Registers hooks to a neural network module to track total operations and parameters."""
         m.register_buffer("total_ops", torch.zeros(1, dtype=torch.float64))
         m.register_buffer("total_params", torch.zeros(1, dtype=torch.float64))
 
@@ -200,6 +204,7 @@ def profile(
         model(*inputs)
 
     def dfs_count(module: nn.Module, prefix="\t") -> (int, int):
+        """Recursively counts the total operations and parameters of the given PyTorch module and its submodules."""
         total_ops, total_params = module.total_ops.item(), 0
         ret_dict = {}
         for n, m in module.named_children():

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/rnn_hooks.py

@@ -4,7 +4,7 @@ from torch.nn.utils.rnn import PackedSequence
 
 
 def _count_rnn_cell(input_size, hidden_size, bias=True):
-    # h' = \tanh(W_{ih} x + b_{ih}  +  W_{hh} h + b_{hh})
+    """Calculate the total operations for an RNN cell based on input size, hidden size, and bias configuration."""
     total_ops = hidden_size * (input_size + hidden_size) + hidden_size
     if bias:
         total_ops += hidden_size * 2
@@ -13,6 +13,7 @@ def _count_rnn_cell(input_size, hidden_size, bias=True):
 
 
 def count_rnn_cell(m: nn.RNNCell, x: torch.Tensor, y: torch.Tensor):
+    """Counts RNN cell operations based on input, hidden size, bias, and batch size."""
     total_ops = _count_rnn_cell(m.input_size, m.hidden_size, m.bias)
 
     batch_size = x[0].size(0)
@@ -22,6 +23,7 @@ def count_rnn_cell(m: nn.RNNCell, x: torch.Tensor, y: torch.Tensor):
 
 
 def _count_gru_cell(input_size, hidden_size, bias=True):
+    """Counts the total operations for a GRU cell based on input size, hidden size, and bias."""
     total_ops = 0
     # r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\
     # z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\
@@ -45,6 +47,7 @@ def _count_gru_cell(input_size, hidden_size, bias=True):
 
 
 def count_gru_cell(m: nn.GRUCell, x: torch.Tensor, y: torch.Tensor):
+    """Calculates and updates the total operations for a GRU cell in a mini-batch during inference."""
     total_ops = _count_gru_cell(m.input_size, m.hidden_size, m.bias)
 
     batch_size = x[0].size(0)
@@ -54,6 +57,9 @@ def count_gru_cell(m: nn.GRUCell, x: torch.Tensor, y: torch.Tensor):
 
 
 def _count_lstm_cell(input_size, hidden_size, bias=True):
+    """Calculates the total operations for an LSTM cell during inference given input size, hidden size, and optional
+    bias.
+    """
     total_ops = 0
 
     # i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\
@@ -76,6 +82,9 @@ def _count_lstm_cell(input_size, hidden_size, bias=True):
 
 
 def count_lstm_cell(m: nn.LSTMCell, x: torch.Tensor, y: torch.Tensor):
+    """Count the number of operations for a single LSTM cell in a given batch, updating the model's total operations
+    count.
+    """
     total_ops = _count_lstm_cell(m.input_size, m.hidden_size, m.bias)
 
     batch_size = x[0].size(0)
@@ -85,6 +94,7 @@ def count_lstm_cell(m: nn.LSTMCell, x: torch.Tensor, y: torch.Tensor):
 
 
 def count_rnn(m: nn.RNN, x, y):
+    """Calculate and update the total number of operations for a single RNN cell in a given batch."""
     bias = m.bias
     input_size = m.input_size
     hidden_size = m.hidden_size
@@ -122,6 +132,7 @@ def count_rnn(m: nn.RNN, x, y):
 
 
 def count_gru(m: nn.GRU, x, y):
+    """Calculate the total number of operations for a GRU layer in a neural network model."""
     bias = m.bias
     input_size = m.input_size
     hidden_size = m.hidden_size
@@ -159,6 +170,9 @@ def count_gru(m: nn.GRU, x, y):
 
 
 def count_lstm(m: nn.LSTM, x, y):
+    """Calculate the total operations for LSTM layers in a network, accounting for input size, hidden size, bias, and
+    bidirectionality.
+    """
     bias = m.bias
     input_size = m.input_size
     hidden_size = m.hidden_size

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/utils.py

@@ -6,6 +6,8 @@ COLOR_YELLOW = "93m"
 
 
 def colorful_print(fn_print, color=COLOR_RED):
+    """A decorator to print text in the specified terminal color by wrapping the given print function."""
+
     def actual_call(*args, **kwargs):
         print(f"\033[{color}", end="")
         fn_print(*args, **kwargs)
@@ -29,6 +31,7 @@ prYellow = colorful_print(print, color=COLOR_YELLOW)
 
 
 def clever_format(nums, format="%.2f"):
+    """Formats numerical values into a more readable string with units (K, M, G, T) based on their magnitude."""
     if not isinstance(nums, Iterable):
         nums = [nums]
     clever_nums = []

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/vision/basic_hooks.py

@@ -11,6 +11,7 @@ multiply_adds = 1
 
 
 def count_parameters(m, x, y):
+    """Calculate and update the total number of parameters in a given PyTorch model."""
     total_params = 0
     for p in m.parameters():
         total_params += torch.DoubleTensor([p.numel()])
@@ -18,10 +19,12 @@ def count_parameters(m, x, y):
 
 
 def zero_ops(m, x, y):
+    """Incrementally add the number of zero operations to the model's total operations count."""
     m.total_ops += calculate_zero_ops()
 
 
 def count_convNd(m: _ConvNd, x, y: torch.Tensor):
+    """Calculate and add the number of convolutional operations (FLOPs) to the model's total operations count."""
     x = x[0]
 
     kernel_ops = torch.zeros(m.weight.size()[2:]).numel()  # Kw x Kh
@@ -45,6 +48,7 @@ def count_convNd(m: _ConvNd, x, y: torch.Tensor):
 
 
 def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor):
+    """Calculates the total operations for a convolutional layer and updates the layer's total_ops attribute."""
     x = x[0]
 
     # N x H x W (exclude Cout)
@@ -60,7 +64,9 @@ def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor):
 
 
 def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y):
-    # TODO: add test cases
+    """Calculate and add the FLOPs for a batch normalization layer, considering elementwise operations and possible
+    affine parameters.
+    """
     # https://github.com/Lyken17/pytorch-OpCounter/issues/124
     # y = (x - mean) / sqrt(eps + var) * weight + bias
     x = x[0]
@@ -82,6 +88,7 @@ def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y):
 
 
 def count_prelu(m, x, y):
+    """Calculate and update the total operation counts for a PReLU layer."""
     x = x[0]
 
     nelements = x.numel()
@@ -90,6 +97,7 @@ def count_prelu(m, x, y):
 
 
 def count_relu(m, x, y):
+    """Calculate and update the total operation counts for a ReLU layer."""
     x = x[0]
 
     nelements = x.numel()
@@ -98,6 +106,7 @@ def count_relu(m, x, y):
 
 
 def count_softmax(m, x, y):
+    """Calculate and update the total operation counts for a Softmax layer."""
     x = x[0]
     nfeatures = x.size()[m.dim]
     batch_size = x.numel() // nfeatures
@@ -106,7 +115,7 @@ def count_softmax(m, x, y):
 
 
 def count_avgpool(m, x, y):
-    # total_add = torch.prod(torch.Tensor([m.kernel_size]))
+    """Calculate and update the total operation counts for an AvgPool layer."""
     # total_div = 1
     # kernel_ops = total_add + total_div
     num_elements = y.numel()
@@ -114,6 +123,7 @@ def count_avgpool(m, x, y):
 
 
 def count_adap_avgpool(m, x, y):
+    """Calculate and update the total operation counts for an AdaptiveAvgPool layer."""
     kernel = torch.div(torch.DoubleTensor([*(x[0].shape[2:])]), torch.DoubleTensor([*(y.shape[2:])]))
     total_add = torch.prod(kernel)
     num_elements = y.numel()
@@ -122,6 +132,7 @@ def count_adap_avgpool(m, x, y):
 
 # TODO: verify the accuracy
 def count_upsample(m, x, y):
+    """Update the total operations counter in the given module for supported upsampling modes."""
     if m.mode not in (
         "nearest",
         "linear",
@@ -137,7 +148,9 @@ def count_upsample(m, x, y):
 
 # nn.Linear
 def count_linear(m, x, y):
-    # per output element
+    """Counts total operations for nn.Linear layers by calculating multiplications and additions based on input and
+    output elements.
+    """
     total_mul = m.in_features
     # total_add = m.in_features - 1
     # total_add += 1 if m.bias is not None else 0

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/vision/calc_func.py

@@ -5,6 +5,7 @@ import torch
 
 
 def l_prod(in_list):
+    """Calculate the product of all elements in a list."""
     res = 1
     for _ in in_list:
         res *= _
@@ -12,6 +13,7 @@ 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 += _
@@ -19,6 +21,7 @@ def l_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()])
@@ -26,11 +29,12 @@ def calculate_parameters(param_list):
 
 
 def calculate_zero_ops():
+    """Return a tensor initialized to zero."""
     return torch.DoubleTensor([int(0)])
 
 
 def calculate_conv2d_flops(input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False):
-    # n, out_c, oh, ow = output_size
+    """Calculate FLOPs for a Conv2D layer given input/output sizes, kernel size, groups, and bias flag."""
     # n, in_c, ih, iw = input_size
     # out_c, in_c, kh, kw = kernel_size
     in_c = input_size[1]
@@ -50,16 +54,20 @@ def calculate_norm(input_size):
 
 
 def calculate_relu_flops(input_size):
-    # x[x < 0] = 0
+    """Calculates the FLOPs for a ReLU activation function based on the input size."""
     return 0
 
 
 def calculate_relu(input_size: torch.Tensor):
+    """Convert an input tensor to a DoubleTensor with the same value."""
     warnings.warn("This API is being deprecated")
     return torch.DoubleTensor([int(input_size)])
 
 
 def calculate_softmax(batch_size, nfeatures):
+    """Calculate the number of FLOPs required for a softmax activation function based on batch size and number of
+    features.
+    """
     total_exp = nfeatures
     total_add = nfeatures - 1
     total_div = nfeatures
@@ -68,16 +76,19 @@ def calculate_softmax(batch_size, nfeatures):
 
 
 def calculate_avgpool(input_size):
+    """Calculate the average pooling size given the input size."""
     return torch.DoubleTensor([int(input_size)])
 
 
 def calculate_adaptive_avg(kernel_size, output_size):
+    """Calculate the number of operations for adaptive average pooling given kernel and output sizes."""
     total_div = 1
     kernel_op = kernel_size + total_div
     return torch.DoubleTensor([int(kernel_op * 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
@@ -93,26 +104,32 @@ def calculate_upsample(mode: str, output_size):
 
 
 def calculate_linear(in_feature, num_elements):
+    """Calculate the linear operation count for an input feature and number of elements."""
     return torch.DoubleTensor([int(in_feature * num_elements)])
 
 
 def counter_matmul(input_size, output_size):
+    """Calculate the total number of operations for a matrix multiplication given input and output sizes."""
     input_size = np.array(input_size)
     output_size = np.array(output_size)
     return np.prod(input_size) * output_size[-1]
 
 
 def counter_mul(input_size):
+    """Calculate the total number of operations for a matrix multiplication given input and output sizes."""
     return input_size
 
 
 def counter_pow(input_size):
+    """Calculate the total number of scalar multiplications for a power operation given an input size."""
     return input_size
 
 
 def counter_sqrt(input_size):
+    """Calculate the total number of scalar operations for a square root operation given an input size."""
     return input_size
 
 
 def counter_div(input_size):
+    """Calculate the total number of scalar operations for a division operation given an input size."""
     return input_size

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/thop/vision/onnx_counter.py

@@ -19,6 +19,7 @@ from .calc_func import (
 
 
 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]
@@ -30,6 +31,7 @@ def onnx_counter_matmul(diction, node):
 
 
 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:
@@ -42,7 +44,9 @@ def onnx_counter_add(diction, node):
 
 
 def onnx_counter_conv(diction, node):
-    # print(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
@@ -81,7 +85,7 @@ def onnx_counter_conv(diction, node):
 
 
 def onnx_counter_constant(diction, node):
-    # print("constant",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]
@@ -90,6 +94,7 @@ def onnx_counter_constant(diction, node):
 
 
 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:
@@ -101,6 +106,7 @@ def onnx_counter_mul(diction, node):
 
 
 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]
@@ -109,6 +115,7 @@ def onnx_counter_bn(diction, node):
 
 
 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]
@@ -120,6 +127,9 @@ def onnx_counter_relu(diction, node):
 
 
 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:
@@ -139,6 +149,7 @@ def onnx_counter_reducemean(diction, node):
 
 
 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]
@@ -147,6 +158,7 @@ def onnx_counter_sub(diction, node):
 
 
 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:
@@ -158,6 +170,7 @@ def onnx_counter_pow(diction, node):
 
 
 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]
@@ -166,6 +179,7 @@ def onnx_counter_sqrt(diction, node):
 
 
 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:
@@ -177,6 +191,7 @@ def onnx_counter_div(diction, node):
 
 
 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]
@@ -185,6 +200,7 @@ def onnx_counter_instance(diction, node):
 
 
 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]
@@ -196,7 +212,7 @@ def onnx_counter_softmax(diction, node):
 
 
 def onnx_counter_pad(diction, node):
-    # # TODO add constant name and output real vector
+    """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]]
@@ -210,7 +226,7 @@ def onnx_counter_pad(diction, node):
 
 
 def onnx_counter_averagepool(diction, node):
-    # TODO add support of ceil_mode and floor
+    """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
@@ -240,7 +256,7 @@ def onnx_counter_averagepool(diction, node):
 
 
 def onnx_counter_flatten(diction, node):
-    # print(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
@@ -251,7 +267,7 @@ def onnx_counter_flatten(diction, node):
 
 
 def onnx_counter_gemm(diction, node):
-    # print(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]]
@@ -264,7 +280,7 @@ def onnx_counter_gemm(diction, node):
 
 
 def onnx_counter_maxpool(diction, node):
-    # TODO add support of ceil_mode and floor
+    """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]
@@ -295,6 +311,7 @@ def onnx_counter_maxpool(diction, node):
 
 
 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]]
@@ -303,7 +320,7 @@ def onnx_counter_globalaveragepool(diction, node):
 
 
 def onnx_counter_concat(diction, node):
-    # print(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]]
@@ -317,6 +334,9 @@ def onnx_counter_concat(diction, node):
 
 
 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]]

{ultralytics_thop-0.0.1 → ultralytics_thop-0.0.2}/ultralytics_thop.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics-thop
-Version: 0.0.1
+Version: 0.0.2
 Summary: A tool to count the FLOPs of PyTorch model.
 Author-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>
 Maintainer-email: Ligeng Zhu <ligeng.zhu+github@gmail.com>