Misidentification of recursiveness (?)

[kswannet]

When using more than one nn.sequential modules, and they both use the same activation functions defined under init, torchinfo splits the single nn.sequential in separate ones at each activation function call.

For example:

class TestNetwork(nn.Module):
    def __init__(self, *, input_dim, h_dim, output_dim, actFun=nn.LeakyReLU(), dropRate=0.3):

        super().__init__()
        self.input_dim = input_dim
        self.h_dim = h_dim
        self.output_dim = output_dim
        self.actFun = actFun
        self.dropRate = dropRate

        layer_dim = [input_dim] + h_dim + [output_dim]   # lengths of all layers

        self.firstNetwork = nn.Sequential()
        for i in range(len(h_dim)+1):
            self.firstNetwork.add_module(f"fcFirst_{i}", nn.Linear(layer_dim[i], layer_dim[i+1]))
            self.firstNetwork.add_module(f"actFirst_{i}", self.actFun)
            self.firstNetwork.add_module(f"dropFirst_{i}", nn.Dropout(dropRate))

        self.secondNetwork = nn.Sequential()
        for i in range(len(h_dim)+1):
            self.secondNetwork.add_module(f"fcSecnd_{i}", nn.Linear(layer_dim[i], layer_dim[i+1]))
            self.secondNetwork.add_module(f"actSecnd_{i}", self.actFun)
            self.secondNetwork.add_module(f"dropSecnd_{i}", nn.Dropout(dropRate))
        del self.secondNetwork[-1]            # remove last dropout layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.firstNetwork(x)

if __name__ == "__main__":
    net = TestNetwork(
        input_dim=199,
        h_dim=[64, 32],
        output_dim=100,
        )
    print(str(summary(net, input_size=[16, 199])))

results in the following summary:

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
TestNetwork                              [16, 100]                 12,800
├─Sequential: 1-1                        [16, 100]                 5,380
│    └─Linear: 2-1                       [16, 64]                  12,800
├─Sequential: 1-6                        --                        (recursive)
│    └─LeakyReLU: 2-2                    [16, 64]                  --
├─Sequential: 1-7                        --                        (recursive)
│    └─Dropout: 2-3                      [16, 64]                  --
│    └─Linear: 2-4                       [16, 32]                  2,080
├─Sequential: 1-6                        --                        (recursive)
│    └─LeakyReLU: 2-5                    [16, 32]                  --
├─Sequential: 1-7                        --                        (recursive)
│    └─Dropout: 2-6                      [16, 32]                  --
│    └─Linear: 2-7                       [16, 100]                 3,300
├─Sequential: 1-6                        --                        (recursive)
│    └─LeakyReLU: 2-8                    [16, 100]                 --
├─Sequential: 1-7                        --                        (recursive)
│    └─Dropout: 2-9                      [16, 100]                 --
==========================================================================================

Even though the secondNetwork is unused, changing one of the self.actFun calls to e.g. nn.LeakyReLU fixes the problem:

class TestNetwork(nn.Module):
    def __init__(self, *, input_dim, h_dim, output_dim, actFun=nn.LeakyReLU(), dropRate=0.3):

        super().__init__()
        self.input_dim = input_dim
        self.h_dim = h_dim
        self.output_dim = output_dim
        self.actFun = actFun
        self.dropRate = dropRate

        layer_dim = [input_dim] + h_dim + [output_dim]   # lengths of all layers

        self.firstNetwork = nn.Sequential()
        for i in range(len(h_dim)+1):
            self.firstNetwork.add_module(f"fcFirst_{i}", nn.Linear(layer_dim[i], layer_dim[i+1]))
            self.firstNetwork.add_module(f"actFirst_{i}", self.actFun)
            self.firstNetwork.add_module(f"dropFirst_{i}", nn.Dropout(dropRate))

        self.secondNetwork = nn.Sequential()
        for i in range(len(h_dim)+1):
            self.secondNetwork.add_module(f"fcSecnd_{i}", nn.Linear(layer_dim[i], layer_dim[i+1]))
            self.secondNetwork.add_module(f"actSecnd_{i}", nn.LeakyReLU())  # Change activation function call
            self.secondNetwork.add_module(f"dropSecnd_{i}", nn.Dropout(dropRate))
        del self.secondNetwork[-1]            # remove last dropout layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.firstNetwork(x)

if __name__ == "__main__":
    net = TestNetwork(
        input_dim=199,
        h_dim=[64, 32],
        output_dim=100,
        )
    print(str(summary(net, input_size=[16, 199])))

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
TestNetwork                              [16, 100]                 18,180
├─Sequential: 1-1                        [16, 100]                 --
│    └─Linear: 2-1                       [16, 64]                  12,800
│    └─LeakyReLU: 2-2                    [16, 64]                  --
│    └─Dropout: 2-3                      [16, 64]                  --
│    └─Linear: 2-4                       [16, 32]                  2,080
│    └─LeakyReLU: 2-5                    [16, 32]                  --
│    └─Dropout: 2-6                      [16, 32]                  --
│    └─Linear: 2-7                       [16, 100]                 3,300
│    └─LeakyReLU: 2-8                    [16, 100]                 --
│    └─Dropout: 2-9                      [16, 100]                 --
==========================================================================================

Is this a torchinfo problem or am I maybe doing something wrong here?

[mykappa]

I have exactly the same problem. With more complex models, the output becomes completely confusing and not well comprehensible.

Update: I've found that adding the the option "hide_recursive_layers" to the row_settings did improve the output a lot for my case.