Add support for pruned models (#103)

* Add support for pruned models

According to the [pytorch documentation on
pruning](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html),
the original parameter is replaced with one ending with `_orig` and
a new buffer ending with `_mask`. The mask contains 0s and 1s based on
which the correct parameters are chosen.

All instances of `param.nelements()` have been replaced by a variable
`cur_params` whose value is set based on whether it is a masked model or
not. To keep consistency with the rest of the code base, the `_orig` is
removed from the `name` variable right after the calculation of
`cur_params`.

* Add tests for pruning

Author: MajorCarrot

tests/test_output/pruning.out

@@ -0,0 +1,20 @@
+==========================================================================================
+Layer (type:depth-idx)                   Output Shape              Param #
+==========================================================================================
+SingleInputNet                           --                        --
+├─Conv2d: 1-1                            [16, 10, 24, 24]          135
+├─Conv2d: 1-2                            [16, 20, 8, 8]            2,520
+├─Dropout2d: 1-3                         [16, 20, 8, 8]            --
+├─Linear: 1-4                            [16, 50]                  8,050
+├─Linear: 1-5                            [16, 10]                  260
+==========================================================================================
+Total params: 10,965
+Trainable params: 10,965
+Non-trainable params: 0
+Total mult-adds (M): 3.96
+==========================================================================================
+Input size (MB): 0.05
+Forward/backward pass size (MB): 0.91
+Params size (MB): 0.04
+Estimated Total Size (MB): 1.00
+==========================================================================================

tests/torchinfo_test.py

@@ -3,6 +3,7 @@
 import torch
 import torchvision  # type: ignore[import]
 from torch import nn
+from torch.nn.utils import prune
 
 from tests.conftest import verify_output_str
 from tests.fixtures.genotype import GenotypeNetwork  # type: ignore[attr-defined]
@@ -162,6 +163,19 @@ def test_resnet152() -> None:
     summary(model, (1, 3, 224, 224), depth=3)
 
 
+def test_pruning() -> None:
+    model = SingleInputNet()
+    for module in model.modules():
+        if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
+            prune.l1_unstructured(  # type: ignore[no-untyped-call]
+                module, "weight", 0.5
+            )
+    results = summary(model, input_size=(16, 1, 28, 28))
+
+    assert results.total_params == 10965
+    assert results.total_mult_adds == 3957600
+
+
 def test_dict_input() -> None:
     # TODO: expand this test to handle intermediate dict layers.
     model = MultipleInputNetDifferentDtypes()

torchinfo/layer_info.py

@@ -12,6 +12,17 @@
 ]
 
 
+def rgetattr(obj: torch.nn.Module, attr: str) -> torch.Tensor:
+    """Get the tensor submodule called attr from obj."""
+    for attr_i in attr.split("."):
+        obj = getattr(obj, attr_i)
+
+    if isinstance(obj, torch.Tensor):
+        return obj
+    else:
+        raise AttributeError(f"{attr} is not a tensor")
+
+
 class LayerInfo:
     """Class that holds information about a layer module."""
 
@@ -116,16 +127,31 @@ def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
                 layer_name += f"-{self.depth_index}"
         return layer_name
 
+    def __get_cur_params(self, name: str, param: torch.Tensor) -> tuple[int, str]:
+        """
+        Get count of number of params, accounting for mask
+        """
+        # Masked models save the parameter with the name "_orig" added
+        # They have a buffer ending with "_mask" which has only 0s and 1s
+        if name[-4:] == "orig":
+            # If a mask exists, the sum of 1s in mask is number of params
+            # Remove "_orig" for better readability and integration
+            return int(torch.sum(rgetattr(self.module, f"{name[:-4]}mask"))), name[:-5]
+        else:
+            return param.nelement(), name
+
     def calculate_num_params(self) -> None:
         """
         Set num_params, trainable, inner_layers, and kernel_size
         using the module's parameters.
         """
         name = ""
         for name, param in self.module.named_parameters():
-            self.num_params += param.nelement()
+            cur_params, name = self.__get_cur_params(name, param)
+
+            self.num_params += cur_params
             if param.requires_grad:
-                self.trainable_params += param.nelement()
+                self.trainable_params += cur_params
 
             ksize = list(param.size())
             if name == "weight":
@@ -137,7 +163,7 @@ def calculate_num_params(self) -> None:
             # RNN modules have inner weights such as weight_ih_l0
             self.inner_layers[name] = {
                 "kernel_size": str(ksize),
-                "num_params": f"├─{param.nelement():,}",
+                "num_params": f"├─{cur_params:,}",
             }
         if self.inner_layers:
             self.inner_layers[name][
@@ -153,18 +179,18 @@ def calculate_macs(self) -> None:
         i.e., taking the batch-dimension into account.
         """
         for name, param in self.module.named_parameters():
+            cur_params, name = self.__get_cur_params(name, param)
             if name in ("weight", "bias"):
                 # ignore C when calculating Mult-Adds in ConvNd
                 if "Conv" in self.class_name:
                     self.macs += int(
-                        param.nelement()
-                        * prod(self.output_size[:1] + self.output_size[2:])
+                        cur_params * prod(self.output_size[:1] + self.output_size[2:])
                     )
                 else:
-                    self.macs += self.output_size[0] * param.nelement()
+                    self.macs += self.output_size[0] * cur_params
             # RNN modules have inner weights such as weight_ih_l0
             elif "weight" in name or "bias" in name:
-                self.macs += prod(self.output_size[:2]) * param.nelement()
+                self.macs += prod(self.output_size[:2]) * cur_params
 
     def check_recursive(self, summary_list: list[LayerInfo]) -> None:
         """