diff --git a/torchsummary/tests/unit_tests/torchsummary_test.py b/torchsummary/tests/unit_tests/torchsummary_test.py index ec4b33e..2d9c09b 100644 --- a/torchsummary/tests/unit_tests/torchsummary_test.py +++ b/torchsummary/tests/unit_tests/torchsummary_test.py @@ -1,10 +1,18 @@ import unittest -from torchsummary import summary, summary_string -from torchsummary.tests.test_models.test_model import SingleInputNet, MultipleInputNet, MultipleInputNetDifferentDtypes + import torch +import torch.nn.utils.parametrize as parametrize + +from torchsummary import summary, summary_string +from torchsummary.tests.test_models.test_model import ( + MultipleInputNet, + MultipleInputNetDifferentDtypes, + SingleInputNet, +) gpu_if_available = "cuda:0" if torch.cuda.is_available() else "cpu" + class torchsummaryTests(unittest.TestCase): def test_single_input(self): model = SingleInputNet() @@ -17,8 +25,7 @@ def test_multiple_input(self): model = MultipleInputNet() input1 = (1, 300) input2 = (1, 300) - total_params, trainable_params = summary( - model, [input1, input2], device="cpu") + total_params, trainable_params = summary(model, [input1, input2], device="cpu") self.assertEqual(total_params, 31120) self.assertEqual(trainable_params, 31120) @@ -44,7 +51,8 @@ def test_multiple_input_types(self): input2 = (1, 300) dtypes = [torch.FloatTensor, torch.LongTensor] total_params, trainable_params = summary( - model, [input1, input2], device="cpu", dtypes=dtypes) + model, [input1, input2], device="cpu", dtypes=dtypes + ) self.assertEqual(total_params, 31120) self.assertEqual(trainable_params, 31120) @@ -54,11 +62,79 @@ def test_single_input(self): model = SingleInputNet() input = (1, 28, 28) result, (total_params, trainable_params) = summary_string( - model, input, device="cpu") + model, input, device="cpu" + ) self.assertEqual(type(result), str) self.assertEqual(total_params, 21840) self.assertEqual(trainable_params, 21840) + def test_frozen_bias_accounting(self): + model = torch.nn.Linear(2, 5) + model.bias.requires_grad = False + _, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + self.assertEqual(total_params, 15) + self.assertEqual(trainable_params, 10) + + def test_non_weight_parameter_is_counted(self): + class ModelWithExtraParameter(torch.nn.Module): + def __init__(self): + super(ModelWithExtraParameter, self).__init__() + self.linear = torch.nn.Linear(2, 2, bias=False) + self.extra = torch.nn.Parameter(torch.ones(3)) + + def forward(self, x): + return self.linear(x) * self.extra[:2] + + model = ModelWithExtraParameter() + _, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + self.assertEqual(total_params, 7) + self.assertEqual(trainable_params, 7) + + def test_parametrization_submodules_are_filtered(self): + class IdentityParametrization(torch.nn.Module): + def forward(self, x): + return x + + model = torch.nn.Linear(2, 2) + parametrize.register_parametrization(model, "weight", IdentityParametrization()) + result, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + + self.assertEqual(total_params, 6) + self.assertEqual(trainable_params, 6) + self.assertNotIn("IdentityParametrization", result) + self.assertNotIn("Displayed (effective) params", result) + self.assertNotIn("Buffer params", result) + + def test_parametrization_owned_parameters_are_counted_on_parent_layer(self): + class ParametrizationWithAuxParameter(torch.nn.Module): + def __init__(self): + super(ParametrizationWithAuxParameter, self).__init__() + self.aux = torch.nn.Parameter(torch.ones(3), requires_grad=False) + + def forward(self, x): + return x + + model = torch.nn.Linear(2, 2) + parametrize.register_parametrization( + model, "weight", ParametrizationWithAuxParameter() + ) + result, (total_params, trainable_params) = summary_string( + model, (1, 2), device="cpu" + ) + + self.assertEqual(total_params, 9) + self.assertEqual(trainable_params, 6) + param_line = next( + line for line in result.splitlines() if "ParametrizedLinear-1" in line + ) + self.assertTrue(param_line.rstrip().endswith("9")) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main(buffer=True) diff --git a/torchsummary/torchsummary.py b/torchsummary/torchsummary.py index 1ed065f..6f26dc6 100644 --- a/torchsummary/torchsummary.py +++ b/torchsummary/torchsummary.py @@ -1,54 +1,125 @@ +from collections import OrderedDict +from typing import Any, List + +import numpy as np import torch import torch.nn as nn -from torch.autograd import Variable -from collections import OrderedDict -import numpy as np +def extract_shapes(x: Any) -> List[List[int]]: + """Extract all tensor shapes from nested structures. + + Args: + x (Any): Input which can be a tensor, list, tuple, or dict. + + Returns: + List[List[int]]: A list of shapes, where each shape is represented as a list of integers. + """ + shapes = [] + + if torch.is_tensor(x): + shapes.append(list(x.size())) + elif isinstance(x, dict): + for v in x.values(): + shapes.extend(extract_shapes(v)) + elif isinstance(x, (list, tuple)): + for v in x: + shapes.extend(extract_shapes(v)) + + return shapes + + +def format_shapes( + layer_name: str, shapes: list, nb_params: torch.Tensor, width: int +) -> str: + """Generate formatted string for layer shapes. + + Args: + layer_name (str): _name of the layer + shapes (list): _list of shapes_ + nb_params (torch.Tensor): _number of parameters_ + width (int): _maximum width of layer names for dynamic formatting_ + + Returns: + str: Formatted string representing the layer, its shapes and number of parameters. + """ + lines = [] + for i, s in enumerate(shapes): + if i == 0: + line = "{:>{w1}} {:>25} {:>15}".format( + layer_name, + "[" + str(s) + ("," if len(shapes) > 1 else "]"), + nb_params, + w1=width, + ) + elif i == len(shapes) - 1: + line = "{:>{w1}} {:>25} {:>15}".format("", str(s) + "]", "", w1=width) + else: + line = "{:>{w1}} {:>25} {:>15}".format("", str(s) + ",", "", w1=width) + lines.append(line) + return "\n".join(lines) + + +def should_register_hook(module_name: str, module: nn.Module) -> bool: + name_parts = module_name.split(".") if module_name else [] + return ( + not isinstance(module, nn.Sequential) + and not isinstance(module, nn.ModuleList) + and "parametrizations" not in name_parts + ) + + +def parameter_owner_name(parameter_name: str) -> str: + if parameter_name.startswith("parametrizations."): + return "" + if ".parametrizations." in parameter_name: + return parameter_name.split(".parametrizations.", 1)[0] + if "." in parameter_name: + return parameter_name.rsplit(".", 1)[0] + return "" + + +def get_module_param_counts(model: nn.Module) -> dict: + module_param_counts = {} + for parameter_name, parameter in model.named_parameters(): + owner_name = parameter_owner_name(parameter_name) + module_param_counts[owner_name] = ( + module_param_counts.get(owner_name, 0) + int(parameter.numel()) + ) + return module_param_counts -def summary(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): - result, params_info = summary_string( - model, input_size, batch_size, device, dtypes) + +def summary( + model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None +): + result, params_info = summary_string(model, input_size, batch_size, device, dtypes) print(result) return params_info -def summary_string(model, input_size, batch_size=-1, device=torch.device('cuda:0'), dtypes=None): - if dtypes == None: - dtypes = [torch.FloatTensor]*len(input_size) +def summary_string( + model, input_size, batch_size=-1, device=torch.device("cuda:0"), dtypes=None +): + if dtypes is None: + dtypes = [torch.FloatTensor] * len(input_size) - summary_str = '' + summary_str = "" + module_param_counts = get_module_param_counts(model) - def register_hook(module): + def register_hook(module_name, module): def hook(module, input, output): class_name = str(module.__class__).split(".")[-1].split("'")[0] module_idx = len(summary) m_key = "%s-%i" % (class_name, module_idx + 1) summary[m_key] = OrderedDict() - summary[m_key]["input_shape"] = list(input[0].size()) - summary[m_key]["input_shape"][0] = batch_size - if isinstance(output, (list, tuple)): - summary[m_key]["output_shape"] = [ - [-1] + list(o.size())[1:] for o in output - ] - else: - summary[m_key]["output_shape"] = list(output.size()) - summary[m_key]["output_shape"][0] = batch_size - - params = 0 - if hasattr(module, "weight") and hasattr(module.weight, "size"): - params += torch.prod(torch.LongTensor(list(module.weight.size()))) - summary[m_key]["trainable"] = module.weight.requires_grad - if hasattr(module, "bias") and hasattr(module.bias, "size"): - params += torch.prod(torch.LongTensor(list(module.bias.size()))) - summary[m_key]["nb_params"] = params - - if ( - not isinstance(module, nn.Sequential) - and not isinstance(module, nn.ModuleList) - ): + summary[m_key]["input_shape"] = extract_shapes(input) + summary[m_key]["output_shape"] = extract_shapes(output) + + summary[m_key]["nb_params"] = module_param_counts.get(module_name, 0) + + if should_register_hook(module_name, module): hooks.append(module.register_forward_hook(hook)) # multiple inputs to the network @@ -56,15 +127,18 @@ def hook(module, input, output): input_size = [input_size] # batch_size of 2 for batchnorm - x = [torch.rand(2, *in_size).type(dtype).to(device=device) - for in_size, dtype in zip(input_size, dtypes)] + x = [ + torch.rand(2, *in_size).type(dtype).to(device=device) + for in_size, dtype in zip(input_size, dtypes) + ] # create properties summary = OrderedDict() hooks = [] # register hook - model.apply(register_hook) + for module_name, module in model.named_modules(): + register_hook(module_name, module) # make a forward pass # print(x.shape) @@ -74,47 +148,48 @@ def hook(module, input, output): for h in hooks: h.remove() - summary_str += "----------------------------------------------------------------" + "\n" - line_new = "{:>20} {:>25} {:>15}".format( - "Layer (type)", "Output Shape", "Param #") + max_layer_len = max(len(layer) for layer in summary) + + summary_str += "-" * (max_layer_len + 43) + "\n" + line_new = "{:>{w1}} {:>25} {:>15}".format( + "Layer (type)", "Output Shape", "Param #", w1=max_layer_len + ) summary_str += line_new + "\n" - summary_str += "================================================================" + "\n" - total_params = 0 + summary_str += "=" * (max_layer_len + 43) + "\n" total_output = 0 - trainable_params = 0 for layer in summary: # input_shape, output_shape, trainable, nb_params - line_new = "{:>20} {:>25} {:>15}".format( - layer, - str(summary[layer]["output_shape"]), - "{0:,}".format(summary[layer]["nb_params"]), + shapes = summary[layer]["output_shape"] + line_new = format_shapes( + layer, shapes, summary[layer]["nb_params"], max_layer_len ) - total_params += summary[layer]["nb_params"] - total_output += np.prod(summary[layer]["output_shape"]) - if "trainable" in summary[layer]: - if summary[layer]["trainable"] == True: - trainable_params += summary[layer]["nb_params"] + for shape in summary[layer]["output_shape"]: + total_output += np.prod(shape) summary_str += line_new + "\n" + params = list(model.parameters()) + total_params = sum(p.numel() for p in params) + trainable_params = sum(p.numel() for p in params if p.requires_grad) + frozen_params = total_params - trainable_params + # assume 4 bytes/number (float on cuda). - total_input_size = abs(np.prod(sum(input_size, ())) - * batch_size * 4. / (1024 ** 2.)) - total_output_size = abs(2. * total_output * 4. / - (1024 ** 2.)) # x2 for gradients - total_params_size = abs(total_params * 4. / (1024 ** 2.)) + total_input_size = abs( + np.prod(sum(input_size, ())) * batch_size * 4.0 / (1024**2.0) + ) + total_output_size = abs(2.0 * total_output * 4.0 / (1024**2.0)) # x2 for gradients + total_params_size = abs(total_params * 4.0 / (1024**2.0)) total_size = total_params_size + total_output_size + total_input_size - summary_str += "================================================================" + "\n" + summary_str += "=" * (max_layer_len + 43) + "\n" summary_str += "Total params: {0:,}".format(total_params) + "\n" summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n" - summary_str += "Non-trainable params: {0:,}".format(total_params - - trainable_params) + "\n" - summary_str += "----------------------------------------------------------------" + "\n" + summary_str += "Non-trainable params: {0:,}".format(frozen_params) + "\n" + summary_str += "-" * (max_layer_len + 43) + "\n" summary_str += "Input size (MB): %0.2f" % total_input_size + "\n" summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n" summary_str += "Params size (MB): %0.2f" % total_params_size + "\n" summary_str += "Estimated Total Size (MB): %0.2f" % total_size + "\n" - summary_str += "----------------------------------------------------------------" + "\n" + summary_str += "-" * (max_layer_len + 43) + "\n" # return summary return summary_str, (total_params, trainable_params)