Allow existing hooks to still work with reused layer variables

Author: TylerYep

tests/fixtures/models.py

@@ -500,6 +500,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return cast(torch.Tensor, self.model(x))
 
 
+class ReuseLinearExtended(nn.Module):
+    """Model that uses a reference to the same Linear layer over and over."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(10, 10)
+        model = []
+        for _ in range(4):
+            model += [self.linear, nn.ReLU(True)]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return cast(torch.Tensor, self.model(x))
+
+
 class ReuseReLU(nn.Module):
     """Model that uses a reference to the same ReLU layer over and over."""
 

tests/test_output/too_many_linear_plus_existing_hooks.out

@@ -0,0 +1,56 @@
+==========================================================================================
+Layer (type:depth-idx)                   Output Shape              Param #
+==========================================================================================
+ReuseLinearExtended                      --                        --
+├─Sequential: 1-1                        [2, 10]                   --
+├─Linear: 1-2                            [2, 10]                   110
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-1                         [2, 10]                   --
+├─Linear: 1-3                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-2                         [2, 10]                   --
+├─Linear: 1-4                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-3                         [2, 10]                   --
+├─Linear: 1-5                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-4                         [2, 10]                   --
+==========================================================================================
+Total params: 110
+Trainable params: 110
+Non-trainable params: 0
+Total mult-adds (M): 0.00
+==========================================================================================
+Input size (MB): 0.00
+Forward/backward pass size (MB): 0.00
+Params size (MB): 0.00
+Estimated Total Size (MB): 0.00
+==========================================================================================
+==========================================================================================
+Layer (type:depth-idx)                   Output Shape              Param #
+==========================================================================================
+ReuseLinearExtended                      --                        --
+├─Sequential: 1-1                        [2, 10]                   --
+├─Linear: 1-2                            [2, 10]                   110
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-1                         [2, 10]                   --
+├─Linear: 1-3                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-2                         [2, 10]                   --
+├─Linear: 1-4                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-3                         [2, 10]                   --
+├─Linear: 1-5                            [2, 10]                   (recursive)
+├─Sequential: 1                          --                        --
+│    └─ReLU: 2-4                         [2, 10]                   --
+==========================================================================================
+Total params: 110
+Trainable params: 110
+Non-trainable params: 0
+Total mult-adds (M): 0.00
+==========================================================================================
+Input size (MB): 0.00
+Forward/backward pass size (MB): 0.00
+Params size (MB): 0.00
+Estimated Total Size (MB): 0.00
+==========================================================================================

tests/torchinfo_test.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import pytest
 import torch
 import torchvision  # type: ignore[import]
@@ -24,6 +26,7 @@
     RecursiveNet,
     ReturnDict,
     ReuseLinear,
+    ReuseLinearExtended,
     ReuseReLU,
     SiameseNets,
     SingleInputNet,
@@ -528,5 +531,32 @@ def test_too_many_linear() -> None:
     summary(net, (2, 10))
 
 
+def test_too_many_linear_plus_existing_hooks() -> None:
+    a, b = False, False
+
+    def pre_hook(module: nn.Module, inputs: Any) -> None:
+        del module, inputs
+        nonlocal a
+        a = True
+
+    def hook(module: nn.Module, inputs: Any, outputs: Any) -> None:
+        del module, inputs, outputs
+        nonlocal b
+        b = True
+
+    net = ReuseLinearExtended()
+    result_1 = summary(net, (2, 10))
+
+    net = ReuseLinearExtended()
+    net.linear.register_forward_pre_hook(pre_hook)
+    net.linear.register_forward_hook(hook)
+
+    result_2 = summary(net, (2, 10))
+
+    assert a is True
+    assert b is True
+    assert str(result_1) == str(result_2)
+
+
 def test_too_many_relus() -> None:
     summary(ReuseReLU(), (4, 4, 64, 64))

torchinfo/torchinfo.py

@@ -251,7 +251,9 @@ def forward_pass(
 
     all_layers: list[LayerInfo] = []
     summary_list: list[LayerInfo] = []
-    hooks: list[RemovableHandle] | None = None if x is None else []
+    hooks: dict[int, tuple[RemovableHandle, RemovableHandle]] | None = (
+        None if x is None else {}
+    )
     named_module = (model_name, model)
     apply_hooks(named_module, model, batch_dim, summary_list, hooks, all_layers)
 
@@ -282,7 +284,8 @@ def forward_pass(
         ) from e
     finally:
         if hooks is not None:
-            for hook in hooks:
+            for pre_hook, hook in hooks.values():
+                pre_hook.remove()
                 hook.remove()
         model.train(saved_model_mode)
 
@@ -473,7 +476,7 @@ def apply_hooks(
     orig_model: nn.Module,
     batch_dim: int | None,
     summary_list: list[LayerInfo],
-    hooks: list[RemovableHandle] | None,
+    hooks: dict[int, tuple[RemovableHandle, RemovableHandle]] | None,
     all_layers: list[LayerInfo],
     curr_depth: int = 0,
     parent_info: LayerInfo | None = None,
@@ -511,10 +514,12 @@ def hook(module: nn.Module, inputs: Any, outputs: Any) -> None:
         if hooks is None or isinstance(module, WRAPPER_MODULES):
             pre_hook(module, None)
         else:
-            if not module._forward_pre_hooks:  # pylint: disable=protected-access
-                hooks.append(module.register_forward_pre_hook(pre_hook))
-            if not module._forward_hooks:  # pylint: disable=protected-access
-                hooks.append(module.register_forward_hook(hook))
+            key = id(module)
+            if key not in hooks:
+                hooks[key] = (
+                    module.register_forward_pre_hook(pre_hook),
+                    module.register_forward_hook(hook),
+                )
 
     # module.named_modules(remove_duplicate=False) doesn't work (infinite recursion).
     for name, mod in module._modules.items():  # pylint: disable=protected-access