Add support for nested/OrderedDict outputs and dynamic layer-aligned formatting for FPN-style backbones

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)

torchsummary/torchsummary.py

@@ -1,70 +1,144 @@
+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
     if isinstance(input_size, tuple):
         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)