From cc279560a41d5081f773b94f43718f14fa28a4f7 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 1 Jul 2024 01:06:05 +0200 Subject: [PATCH] Ultralytics Code Refactor https://ultralytics.com/actions (#32) Co-authored-by: UltralyticsAssistant --- tests/test_conv2d.py | 6 ++---- tests/test_matmul.py | 6 +++--- tests/test_relu.py | 4 +--- tests/test_utils.py | 6 ++---- thop/fx_profile.py | 20 ++++++++++---------- thop/profile.py | 6 ++---- thop/rnn_hooks.py | 22 ++++++++-------------- thop/utils.py | 2 +- thop/vision/basic_hooks.py | 26 +++++++++++--------------- thop/vision/calc_func.py | 36 +++++++++++++++++------------------- 10 files changed, 57 insertions(+), 77 deletions(-) diff --git a/tests/test_conv2d.py b/tests/test_conv2d.py index e058e25..11b000e 100644 --- a/tests/test_conv2d.py +++ b/tests/test_conv2d.py @@ -6,9 +6,7 @@ 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. - """ + """Tests a 2D Conv layer without bias using THOP profiling on predefined input dimensions and 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 @@ -23,7 +21,7 @@ def test_conv2d_no_bias(self): 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.""" + """Tests Conv2D layer with bias, profiling FLOPs and params for specific input dimensions and layer configs.""" 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 diff --git a/tests/test_matmul.py b/tests/test_matmul.py index 83cc361..227b94f 100644 --- a/tests/test_matmul.py +++ b/tests/test_matmul.py @@ -6,7 +6,7 @@ class TestUtils: def test_matmul_case2(self): - """Test matrix multiplication case asserting the FLOPs and parameters of a nn.Linear layer.""" + """Test matrix multiplication case by profiling FLOPs and parameters of a PyTorch 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,7 +14,7 @@ def test_matmul_case2(self): 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.""" + """Tests matrix multiplication to profile FLOPs and parameters of nn.Linear layer using random dimensions.""" for _ in range(10): n, in_c, out_c = torch.randint(1, 500, (3,)).tolist() net = nn.Linear(in_c, out_c) @@ -23,7 +23,7 @@ def test_matmul_case2(self): 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.""" + """Tests FLOPs and parameters for a nn.Linear layer with random dimensions 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),)) diff --git a/tests/test_relu.py b/tests/test_relu.py index 88739e6..d45518e 100644 --- a/tests/test_relu.py +++ b/tests/test_relu.py @@ -6,9 +6,7 @@ class TestUtils: def test_relu(self): - """Tests the ReLU activation function to ensure it has zero FLOPs and checks parameter count using THOP - profiling. - """ + """Tests ReLU activation ensuring zero FLOPs and displays parameter count using THOP profiling.""" n, in_c, _out_c = 1, 100, 200 net = nn.ReLU() flops, params = profile(net, inputs=(torch.randn(n, in_c),)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 0358e7e..28daa68 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,16 +3,14 @@ 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.""" + """Tests that clever_format returns a string like '1.00B' for the given number and format 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. - """ + """Verifies clever_format formats a list of numbers to strings, e.g., '[1, 2]' to '[1.00B, 2.00B]'.""" nums = [1, 2] format = "%.2f" clever_nums = utils.clever_format(nums, format) diff --git a/thop/fx_profile.py b/thop/fx_profile.py index 6f88f03..48e8ab8 100644 --- a/thop/fx_profile.py +++ b/thop/fx_profile.py @@ -13,7 +13,7 @@ def count_clamp(input_shapes, output_shapes): - """Ensures proper array sizes for tensors by clamping input and output shapes.""" + """Ensures tensor array sizes are appropriate by clamping specified input and output shapes.""" return 0 @@ -23,7 +23,7 @@ def count_mul(input_shapes, output_shapes): def count_matmul(input_shapes, output_shapes): - """Calculates the total number of operations for a matrix multiplication given input and output shapes.""" + """Calculates matrix multiplication ops based on input and output tensor shapes for performance profiling.""" in_shape = input_shapes[0] out_shape = output_shapes[0] in_features = in_shape[-1] @@ -32,7 +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.""" + """Calculates the total FLOPs for a linear layer, including bias operations if specified.""" flops = count_matmul(input_shapes, output_shapes) if "bias" in kwargs: flops += output_shapes[0].numel() @@ -43,7 +43,9 @@ def count_fn_linear(input_shapes, output_shapes, *args, **kwargs): def count_fn_conv2d(input_shapes, output_shapes, *args, **kwargs): - """Calculates total operations (FLOPs) for a 2D convolutional layer given input and output shapes.""" + """Calculates total operations (FLOPs) for a 2D conv layer based on input and output shapes using + `calculate_conv`. + """ inputs, weight, bias, stride, padding, dilation, groups = args if len(input_shapes) == 2: x_shape, k_shape = input_shapes @@ -65,12 +67,12 @@ def count_nn_linear(module: nn.Module, 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.""" + """Returns 0 for a neural network module, input shapes, and output shapes in PyTorch.""" 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.""" + """Calculates FLOPs for a 2D Conv2D layer in an nn.Module using input and output shapes.""" bias_op = 1 if module.bias is not None else 0 out_shape = output_shapes[0] @@ -82,7 +84,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.""" + """Calculate FLOPs for an nn.BatchNorm2d layer based on the given output shape.""" assert len(output_shapes) == 1, "nn.BatchNorm2d should only have one output" y = output_shapes[0] return 2 * y.numel() @@ -127,9 +129,7 @@ def null_print(*args, **kwargs): 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. - """ + """Profiles nn.Module for total FLOPs per operation and prints detailed nodes if verbose.""" gm: torch.fx.GraphModule = symbolic_trace(mod) ShapeProp(gm).propagate(input) diff --git a/thop/profile.py b/thop/profile.py index 4228ca7..8333fcc 100644 --- a/thop/profile.py +++ b/thop/profile.py @@ -55,9 +55,7 @@ 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. - """ + """Profiles a PyTorch model's operations and parameters, applying either custom or default hooks.""" handler_collection = [] types_collection = set() if custom_ops is None: @@ -145,7 +143,7 @@ def profile( ret_layer_info=False, report_missing=False, ): - """Profiles a PyTorch model, returning total operations and parameters, with optional layer-wise details.""" + """Profiles a PyTorch model, returning total operations, parameters, and optionally layer-wise details.""" handler_collection = {} types_collection = set() if custom_ops is None: diff --git a/thop/rnn_hooks.py b/thop/rnn_hooks.py index d87186f..dbf4a01 100644 --- a/thop/rnn_hooks.py +++ b/thop/rnn_hooks.py @@ -4,7 +4,7 @@ def _count_rnn_cell(input_size, hidden_size, bias=True): - """Calculate the total operations for an RNN cell based on input size, hidden size, and bias configuration.""" + """Calculate the total operations for an RNN cell given input size, hidden size, and optional bias.""" total_ops = hidden_size * (input_size + hidden_size) + hidden_size if bias: total_ops += hidden_size * 2 @@ -13,7 +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.""" + """Counts the total RNN cell operations based on input tensor, 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) @@ -23,7 +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.""" + """Counts the total operations for a GRU cell based on input size, hidden size, and bias configuration.""" 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}) \\ @@ -57,9 +57,7 @@ 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. - """ + """Counts LSTM cell operations during inference based on input size, hidden size, and bias configuration.""" total_ops = 0 # i = \sigma(W_{ii} x + b_{ii} + W_{hi} h + b_{hi}) \\ @@ -82,9 +80,7 @@ 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. - """ + """Counts and updates the total operations for an LSTM cell in a mini-batch during inference.""" total_ops = _count_lstm_cell(m.input_size, m.hidden_size, m.bias) batch_size = x[0].size(0) @@ -94,7 +90,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.""" + """Calculate and update the total number of operations for each RNN cell in a given batch.""" bias = m.bias input_size = m.input_size hidden_size = m.hidden_size @@ -131,7 +127,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.""" + """Calculates total operations for a GRU layer, updating the model's operation count based on batch size.""" bias = m.bias input_size = m.input_size hidden_size = m.hidden_size @@ -168,9 +164,7 @@ 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. - """ + """Calculate total operations for LSTM layers, including bidirectional, updating model's total operations.""" bias = m.bias input_size = m.input_size hidden_size = m.hidden_size diff --git a/thop/utils.py b/thop/utils.py index 04114d1..845b54f 100644 --- a/thop/utils.py +++ b/thop/utils.py @@ -22,7 +22,7 @@ def actual_call(*args, **kwargs): def clever_format(nums, format="%.2f"): - """Formats numerical values into a more readable string with units (K, M, G, T) based on their magnitude.""" + """Formats numbers into human-readable strings with units (K for thousand, M for million, etc.).""" if not isinstance(nums, Iterable): nums = [nums] clever_nums = [] diff --git a/thop/vision/basic_hooks.py b/thop/vision/basic_hooks.py index f4dfce3..fb8602d 100644 --- a/thop/vision/basic_hooks.py +++ b/thop/vision/basic_hooks.py @@ -9,17 +9,17 @@ def count_parameters(m, x, y): - """Calculate and update the total number of parameters in a given PyTorch model.""" + """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 the number of zero operations to the model's total operations count.""" + """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) to the model's total operations count.""" + """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( @@ -40,7 +40,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.""" + """Calculates and updates total operations (FLOPs) for a convolutional layer in a PyTorch model.""" x = x[0] # N x H x W (exclude Cout) @@ -56,9 +56,7 @@ def count_convNd_ver2(m: _ConvNd, x, y: torch.Tensor): def count_normalization(m: nn.modules.batchnorm._BatchNorm, x, y): - """Calculate and add the FLOPs for a batch normalization layer, considering elementwise operations and possible - affine parameters. - """ + """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] @@ -80,7 +78,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.""" + """Calculate and update the total operation counts for a PReLU layer using input element number.""" x = x[0] nelements = x.numel() @@ -95,7 +93,7 @@ def count_relu(m, x, y): def count_softmax(m, x, y): - """Calculate and update the total operation counts for a Softmax layer.""" + """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 @@ -104,7 +102,7 @@ def count_softmax(m, x, y): def count_avgpool(m, x, y): - """Calculate and update the total operation counts for an AvgPool layer.""" + """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() @@ -112,7 +110,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.""" + """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() @@ -121,7 +119,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.""" + """Update total operations counter for upsampling layers based on the mode used.""" if m.mode not in ( "nearest", "linear", @@ -137,9 +135,7 @@ def count_upsample(m, x, y): # nn.Linear def count_linear(m, x, y): - """Counts total operations for nn.Linear layers by calculating multiplications and additions based on input and - output elements. - """ + """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 diff --git a/thop/vision/calc_func.py b/thop/vision/calc_func.py index b91648a..9e534ac 100644 --- a/thop/vision/calc_func.py +++ b/thop/vision/calc_func.py @@ -5,7 +5,7 @@ def l_prod(in_list): - """Calculate the product of all elements in a list.""" + """Compute the product of all elements in the input list.""" res = 1 for _ in in_list: res *= _ @@ -13,22 +13,22 @@ def l_prod(in_list): def l_sum(in_list): - """Calculate the sum of all elements in a 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.""" + """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(): - """Return a tensor initialized to zero.""" + """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 given input/output sizes, kernel size, groups, and bias flag.""" + """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] @@ -43,25 +43,23 @@ def calculate_conv(bias, kernel_size, output_size, in_channel, group): def calculate_norm(input_size): - """Input is a number not a array or tensor.""" + """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 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.""" + """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): - """Calculate the number of FLOPs required for a softmax activation function based on batch size and number of - features. - """ + """Compute FLOPs for a softmax activation given batch size and feature count.""" total_exp = nfeatures total_add = nfeatures - 1 total_div = nfeatures @@ -70,19 +68,19 @@ def calculate_softmax(batch_size, nfeatures): def calculate_avgpool(input_size): - """Calculate the average pooling size given the 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 the number of operations for adaptive average pooling given kernel and output sizes.""" + """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 number of operations for upsample methods given the mode and 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 @@ -96,29 +94,29 @@ 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.""" + """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 a matrix multiplication given input and output sizes.""" + """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 a matrix multiplication given input and output sizes.""" + """Calculate the total number of operations for element-wise multiplication given the input size.""" return input_size def counter_pow(input_size): - """Calculate the total number of scalar multiplications for a power operation given an 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 for a square root operation given an input size.""" + """Calculate the total number of scalar operations required for a square root operation given an input size.""" return input_size