Fix reuse of layer variables bug

Author: TylerYep

tests/fixtures/models.py

@@ -16,7 +16,7 @@ class IdentityModel(nn.Module):
 
     def __init__(self) -> None:
         super().__init__()
-        self.identity = nn.Identity()  # type: ignore[no-untyped-call] # noqa
+        self.identity = nn.Identity()  # type: ignore[no-untyped-call]
 
     def forward(self, x: Any) -> Any:
         return self.identity(x)
@@ -483,3 +483,43 @@ def __init__(self) -> None:
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.w * x
+
+
+class ReuseLinear(nn.Module):
+    """Model that uses a reference to the same Linear layer over and over."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        linear = nn.Linear(10, 10)
+        model = []
+        for _ in range(4):
+            model += [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."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        activation = nn.ReLU(True)
+        model = [
+            nn.ReflectionPad2d(3),
+            nn.Conv2d(4, 1, kernel_size=1, padding=0),
+            nn.BatchNorm2d(1),  # type: ignore[no-untyped-call]
+            activation,
+        ]
+        for i in range(3):
+            mult = 2 ** i
+            model += [
+                nn.Conv2d(mult, mult * 2, kernel_size=1, stride=2, padding=1),
+                nn.BatchNorm2d(mult * 2),  # type: ignore[no-untyped-call]
+                activation,
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return cast(torch.Tensor, self.model(x))

tests/test_output/too_many_linear.out

@@ -0,0 +1,24 @@
+==========================================================================================
+Layer (type:depth-idx)                   Output Shape              Param #
+==========================================================================================
+ReuseLinear                              --                        --
+├─Sequential: 1-1                        [2, 10]                   --
+│    └─Linear: 2-1                       [2, 10]                   110
+│    └─ReLU: 2-2                         [2, 10]                   --
+│    └─Linear: 2-3                       [2, 10]                   (recursive)
+│    └─ReLU: 2-4                         [2, 10]                   --
+│    └─Linear: 2-5                       [2, 10]                   (recursive)
+│    └─ReLU: 2-6                         [2, 10]                   --
+│    └─Linear: 2-7                       [2, 10]                   (recursive)
+│    └─ReLU: 2-8                         [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/test_output/too_many_relus.out

@@ -0,0 +1,29 @@
+==========================================================================================
+Layer (type:depth-idx)                   Output Shape              Param #
+==========================================================================================
+ReuseReLU                                --                        --
+├─Sequential: 1-1                        [4, 8, 11, 11]            --
+│    └─ReflectionPad2d: 2-1              [4, 4, 70, 70]            --
+│    └─Conv2d: 2-2                       [4, 1, 70, 70]            5
+│    └─BatchNorm2d: 2-3                  [4, 1, 70, 70]            2
+│    └─ReLU: 2-4                         [4, 1, 70, 70]            --
+│    └─Conv2d: 2-5                       [4, 2, 36, 36]            4
+│    └─BatchNorm2d: 2-6                  [4, 2, 36, 36]            4
+│    └─ReLU: 2-7                         [4, 2, 36, 36]            --
+│    └─Conv2d: 2-8                       [4, 4, 19, 19]            12
+│    └─BatchNorm2d: 2-9                  [4, 4, 19, 19]            8
+│    └─ReLU: 2-10                        [4, 4, 19, 19]            --
+│    └─Conv2d: 2-11                      [4, 8, 11, 11]            40
+│    └─BatchNorm2d: 2-12                 [4, 8, 11, 11]            16
+│    └─ReLU: 2-13                        [4, 8, 11, 11]            --
+==========================================================================================
+Total params: 91
+Trainable params: 91
+Non-trainable params: 0
+Total mult-adds (M): 0.16
+==========================================================================================
+Input size (MB): 0.26
+Forward/backward pass size (MB): 0.63
+Params size (MB): 0.00
+Estimated Total Size (MB): 0.90
+==========================================================================================

tests/torchinfo_test.py

@@ -23,6 +23,8 @@
     PartialJITModel,
     RecursiveNet,
     ReturnDict,
+    ReuseLinear,
+    ReuseReLU,
     SiameseNets,
     SingleInputNet,
 )
@@ -519,3 +521,12 @@ def test_ascii_only() -> None:
 
 def test_google() -> None:
     summary(torchvision.models.googlenet(), (1, 3, 112, 112), depth=7)
+
+
+def test_too_many_linear() -> None:
+    net = ReuseLinear()
+    summary(net, (2, 10))
+
+
+def test_too_many_relus() -> None:
+    summary(ReuseReLU(), (4, 4, 64, 64))

torchinfo/torchinfo.py

@@ -511,8 +511,10 @@ def hook(module: nn.Module, inputs: Any, outputs: Any) -> None:
         if hooks is None or isinstance(module, WRAPPER_MODULES):
             pre_hook(module, None)
         else:
-            hooks.append(module.register_forward_pre_hook(pre_hook))
-            hooks.append(module.register_forward_hook(hook))
+            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))
 
     # module.named_modules(remove_duplicate=False) doesn't work (infinite recursion).
     for name, mod in module._modules.items():  # pylint: disable=protected-access