ultralytics-thop 2.0.14__py3-none-any.whl

thop/utils.py

@@ -0,0 +1,50 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from collections.abc import Iterable
+
+COLOR_RED = "91m"
+COLOR_GREEN = "92m"
+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)
+        print("\033[00m", end="")
+
+    return actual_call
+
+
+prRed = colorful_print(print, color=COLOR_RED)
+prGreen = colorful_print(print, color=COLOR_GREEN)
+prYellow = colorful_print(print, color=COLOR_YELLOW)
+
+
+def clever_format(nums, format="%.2f"):
+    """Formats numbers into human-readable strings with units (K for thousand, M for million, etc.)."""
+    if not isinstance(nums, Iterable):
+        nums = [nums]
+    clever_nums = []
+
+    for num in nums:
+        if num > 1e12:
+            clever_nums.append(format % (num / 1e12) + "T")
+        elif num > 1e9:
+            clever_nums.append(format % (num / 1e9) + "G")
+        elif num > 1e6:
+            clever_nums.append(format % (num / 1e6) + "M")
+        elif num > 1e3:
+            clever_nums.append(format % (num / 1e3) + "K")
+        else:
+            clever_nums.append(format % num + "B")
+
+    return clever_nums[0] if len(clever_nums) == 1 else (*clever_nums,)
+
+
+if __name__ == "__main__":
+    prRed("hello", "world")
+    prGreen("hello", "world")
+    prYellow("hello", "world")

thop/vision/__init__.py

@@ -0,0 +1 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

thop/vision/basic_hooks.py

@@ -0,0 +1,146 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import logging
+
+import torch.nn as nn
+from torch.nn.modules.conv import _ConvNd
+
+from .calc_func import *
+
+multiply_adds = 1
+
+
+def count_parameters(m, x, y):
+    """Calculate and return the total number of learnable parameters in a given PyTorch model."""
+    m.total_params[0] = calculate_parameters(m.parameters())
+
+
+def zero_ops(m, x, y):
+    """Incrementally add 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) for a ConvNd layer to the model's total ops."""
+    x = x[0]
+
+    m.total_ops += calculate_conv2d_flops(
+        input_size=list(x.shape),
+        output_size=list(y.shape),
+        kernel_size=list(m.weight.shape),
+        groups=m.groups,
+        bias=m.bias,
+    )
+    # N x Cout x H x W x  (Cin x Kw x Kh + bias)
+    # m.total_ops += calculate_conv(
+    #     bias_ops,
+    #     torch.zeros(m.weight.size()[2:]).numel(),
+    #     y.nelement(),
+    #     m.in_channels,
+    #     m.groups,
+    # )
+
+
+def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor):
+    """Calculates and updates total operations (FLOPs) for a convolutional layer in a PyTorch model."""
+    x = x[0]
+
+    # N x H x W (exclude Cout)
+    output_size = torch.zeros(y.size()[:1] + y.size()[2:]).numel()
+    # # Cout x Cin x Kw x Kh
+    # kernel_ops = m.weight.nelement()
+    # if m.bias is not None:
+    #     # Cout x 1
+    #     kernel_ops += + m.bias.nelement()
+    # # x N x H x W x Cout x (Cin x Kw x Kh + bias)
+    # m.total_ops += torch.DoubleTensor([int(output_size * kernel_ops)])
+    m.total_ops += calculate_conv(m.bias.nelement(), m.weight.nelement(), output_size)
+
+
+def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y):
+    """Calculate and add the FLOPs for a batch normalization layer, including elementwise and affine operations."""
+    # https://github.com/Lyken17/pytorch-OpCounter/issues/124
+    # y = (x - mean) / sqrt(eps + var) * weight + bias
+    x = x[0]
+    # bn is by default fused in inference
+    flops = calculate_norm(x.numel())
+    if getattr(m, "affine", False) or getattr(m, "elementwise_affine", False):
+        flops *= 2
+    m.total_ops += flops
+
+
+# def count_layer_norm(m, x, y):
+#     x = x[0]
+#     m.total_ops += calculate_norm(x.numel())
+
+
+# def count_instance_norm(m, x, y):
+#     x = x[0]
+#     m.total_ops += calculate_norm(x.numel())
+
+
+def count_prelu(m, x, y):
+    """Calculate and update the total operation counts for a PReLU layer using input element number."""
+    x = x[0]
+
+    nelements = x.numel()
+    if not m.training:
+        m.total_ops += calculate_relu(nelements)
+
+
+def count_relu(m, x, y):
+    """Calculate and update the total operation counts for a ReLU layer."""
+    x = x[0]
+    m.total_ops += calculate_relu_flops(list(x.shape))
+
+
+def count_softmax(m, x, y):
+    """Calculate and update the total operation counts for a Softmax layer in a PyTorch model."""
+    x = x[0]
+    nfeatures = x.size()[m.dim]
+    batch_size = x.numel() // nfeatures
+
+    m.total_ops += calculate_softmax(batch_size, nfeatures)
+
+
+def count_avgpool(m, x, y):
+    """Calculate and update the total number of operations (FLOPs) for an AvgPool layer based on the output elements."""
+    # total_div = 1
+    # kernel_ops = total_add + total_div
+    num_elements = y.numel()
+    m.total_ops += calculate_avgpool(num_elements)
+
+
+def count_adap_avgpool(m, x, y):
+    """Calculate and update the total operation counts for an AdaptiveAvgPool layer using kernel and element counts."""
+    kernel = torch.div(torch.DoubleTensor([*(x[0].shape[2:])]), torch.DoubleTensor([*(y.shape[2:])]))
+    total_add = torch.prod(kernel)
+    num_elements = y.numel()
+    m.total_ops += calculate_adaptive_avg(total_add, num_elements)
+
+
+# TODO: verify the accuracy
+def count_upsample(m, x, y):
+    """Update total operations counter for upsampling layers based on the mode used."""
+    if m.mode not in (
+        "nearest",
+        "linear",
+        "bilinear",
+        "bicubic",
+    ):  # "trilinear"
+        logging.warning(f"mode {m.mode} is not implemented yet, take it a zero op")
+        m.total_ops += 0
+    else:
+        x = x[0]
+        m.total_ops += calculate_upsample(m.mode, y.nelement())
+
+
+# nn.Linear
+def count_linear(m, x, y):
+    """Counts total operations for nn.Linear layers using input and output element dimensions."""
+    total_mul = m.in_features
+    # total_add = m.in_features - 1
+    # total_add += 1 if m.bias is not None else 0
+    num_elements = y.numel()
+
+    m.total_ops += calculate_linear(total_mul, num_elements)

thop/vision/calc_func.py

@@ -0,0 +1,127 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+import warnings
+
+import numpy as np
+import torch
+
+
+def l_prod(in_list):
+    """Compute the product of all elements in the input list."""
+    res = 1
+    for _ in in_list:
+        res *= _
+    return res
+
+
+def l_sum(in_list):
+    """Calculate the sum of all numerical elements in a list."""
+    return sum(in_list)
+
+
+def calculate_parameters(param_list):
+    """Calculate the total number of parameters in a list of tensors using the product of their shapes."""
+    return sum(torch.DoubleTensor([p.nelement()]) for p in param_list)
+
+
+def calculate_zero_ops():
+    """Initializes and returns a tensor with all elements set to zero."""
+    return torch.DoubleTensor([0])
+
+
+def calculate_conv2d_flops(input_size: list, output_size: list, kernel_size: list, groups: int, bias: bool = False):
+    """Calculate FLOPs for a Conv2D layer using input/output sizes, kernel size, groups, and the bias flag."""
+    # n, in_c, ih, iw = input_size
+    # out_c, in_c, kh, kw = kernel_size
+    in_c = input_size[1]
+    g = groups
+    return l_prod(output_size) * (in_c // g) * l_prod(kernel_size[2:])
+
+
+def calculate_conv(bias, kernel_size, output_size, in_channel, group):
+    """Calculate FLOPs for convolutional layers given bias, kernel size, output size, in_channels, and groups."""
+    warnings.warn("This API is being deprecated.")
+    return torch.DoubleTensor([output_size * (in_channel / group * kernel_size + bias)])
+
+
+def calculate_norm(input_size):
+    """Compute the L2 norm of a tensor or array based on its input size."""
+    return torch.DoubleTensor([2 * input_size])
+
+
+def calculate_relu_flops(input_size):
+    """Calculates the FLOPs for a ReLU activation function based on the input tensor's dimensions."""
+    return 0
+
+
+def calculate_relu(input_size: torch.Tensor):
+    """Convert an input tensor to a DoubleTensor with the same value (deprecated)."""
+    warnings.warn("This API is being deprecated")
+    return torch.DoubleTensor([int(input_size)])
+
+
+def calculate_softmax(batch_size, nfeatures):
+    """Compute FLOPs for a softmax activation given batch size and feature count."""
+    total_exp = nfeatures
+    total_add = nfeatures - 1
+    total_div = nfeatures
+    total_ops = batch_size * (total_exp + total_add + total_div)
+    return torch.DoubleTensor([int(total_ops)])
+
+
+def calculate_avgpool(input_size):
+    """Calculate the average pooling size for a given input tensor."""
+    return torch.DoubleTensor([int(input_size)])
+
+
+def calculate_adaptive_avg(kernel_size, output_size):
+    """Calculate FLOPs for adaptive average pooling given kernel size and output size."""
+    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 operations required for various upsample methods based on mode and output size."""
+    total_ops = output_size
+    if mode == "bicubic":
+        total_ops *= 224 + 35
+    elif mode == "bilinear":
+        total_ops *= 11
+    elif mode == "linear":
+        total_ops *= 5
+    elif mode == "trilinear":
+        total_ops *= 13 * 2 + 5
+    return torch.DoubleTensor([int(total_ops)])
+
+
+def calculate_linear(in_feature, num_elements):
+    """Calculate the linear operation count for given 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 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 element-wise multiplication given the input size."""
+    return input_size
+
+
+def counter_pow(input_size):
+    """Computes the total scalar multiplications required for power operations based on input size."""
+    return input_size
+
+
+def counter_sqrt(input_size):
+    """Calculate the total number of scalar operations required 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