Children modules of Repeated Container Modules

[mert-kurttutan]

Describe the bug

At the moment, I don't think this is a bug, but rather a result of recursive use of the same Container Modules.
When there is a repetition of container modules , the children of container modules will set as many times as the repetition of the container module. Therefore, the final value of children modules of the container module will be set by the final repetition

To Reproduce
Use the following code and summary

import torch
import torch.nn as nn

from torchinfo_local.torchinfo import summary



class RecursiveTest(nn.Module):
    def __init__(self):
        super().__init__()
        self.out_lin0 = nn.Linear(128, 4)

        self.block0 = nn.ModuleDict({"activation": nn.ReLU()})
        # range_1 = range(1,4) 
        range_1 = reversed(range(1,4))
        for i in range_1:
            self.block0.add_module(f"in_lin{i}", nn.Linear(4,4)) 

        self.block1 = nn.ModuleDict()
        for i in range(4, 7):
            self.block1.add_module(f"in_lin{i}", nn.Linear(4, 4))

        self.out_lin7 = nn.Linear(4, 2)

    def forward(self, x):
        x = torch.relu(self.out_lin0(x))

        for i in range(1, 4):
            x = self.block0["activation"](self.block0[f"in_lin{i}"](x))


        for i in range(4, 7):
            x = torch.relu(self.block1[f"in_lin{i}"](x))


        x = torch.relu(self.block0[f"in_lin{1}"](x))
        x = torch.relu(self.out_lin7(x))

        return x


batch_size = 2
data_shape = (128,)
random_data = torch.rand((batch_size, *data_shape))
my_nn = RecursiveTest()


recursive_summary = summary(
    my_nn, 
    input_data=[random_data], 
    row_settings=('depth', 'var_names'),
    device='cpu',
)

==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
RecursiveTest (RecursiveTest)            [2, 2]                    --
├─Linear (out_lin0): 1-1                 [2, 4]                    516
├─ModuleDict (block0): 1-4               --                        (recursive)
│    └─Linear (in_lin1): 2-1             [2, 4]                    20
│    └─ReLU (activation): 2-2            [2, 4]                    --
│    └─Linear (in_lin2): 2-3             [2, 4]                    20
│    └─ReLU (activation): 2-4            [2, 4]                    --
│    └─Linear (in_lin3): 2-5             [2, 4]                    20
│    └─ReLU (activation): 2-6            [2, 4]                    --
├─ModuleDict (block1): 1-3               --                        --
│    └─Linear (in_lin4): 2-7             [2, 4]                    20
│    └─Linear (in_lin5): 2-8             [2, 4]                    20
│    └─Linear (in_lin6): 2-9             [2, 4]                    20
├─ModuleDict (block0): 1-4               --                        (recursive)
│    └─Linear (in_lin1): 2-10            [2, 4]                    (recursive)
├─Linear (out_lin7): 1-5                 [2, 2]                    10
==========================================================================================
Total params: 646
Trainable params: 646
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
==========================================================================================

Expected behavior
This does not affect the results in a negative way since children modules are added to the summary list while traversing. So, it uses the current child.
But after the traversal of entire main module, the end result is determined by the last occurrence of container modules. I think there might a better way to set the children of container modules after traversal. Maybe set the children according definition in __init__

[TylerYep]

I don't completely understand what you are noticing here, maybe annotating the example given would help me understand what the confusion is? It sounds like recursive children modules do not contain the correct children on output?

[mert-kurttutan]

Yeah sure,
In the above example, during traversal by add_container_module function, block0 Moduledict essentially sets its children twice.
First, it stops at

│    └─Linear (in_lin1): 2-1             [2, 4]                    20
│    └─ReLU (activation): 2-2            [2, 4]                    --
│    └─Linear (in_lin2): 2-3             [2, 4]                    20
│    └─ReLU (activation): 2-4            [2, 4]                    --
│    └─Linear (in_lin3): 2-5             [2, 4]                    20
│    └─ReLU (activation): 2-6            [2, 4]                    --

Since all of these modules belongs to block0, the children of block0 is set to [Linear: 1, ModuleDict: 1, Linear: 2, ReLU: 2, Linear: 2, ReLU: 2, Linear: 2, ReLU: 2]. Then, it goes block1 and sets the children of block1. Finally, there is another submodule that belong to block0. Here it stops again. This time, it overrides previous children setting and sets its children to [Linear: 2 (in_lin1)].
So the final choice for the children is determined by the last occurrence.

For clarification, I added updated children to the row at which they are set

==========================================================================================
Layer (type (var_name):depth-idx)        Output Shape              Param #
==========================================================================================
RecursiveTest (RecursiveTest)            [2, 2]                    --
├─Linear (out_lin0): 1-1                 [2, 4]                    516
├─ModuleDict (block0): 1-4               --                        (recursive)     (block0_children=[Linear: 1, ModuleDict: 1, Linear: 2, ReLU: 2, Linear: 2, ReLU: 2, Linear: 2, ReLU: 2]
│    └─Linear (in_lin1): 2-1             [2, 4]                    20
│    └─ReLU (activation): 2-2            [2, 4]                    --
│    └─Linear (in_lin2): 2-3             [2, 4]                    20
│    └─ReLU (activation): 2-4            [2, 4]                    --
│    └─Linear (in_lin3): 2-5             [2, 4]                    20
│    └─ReLU (activation): 2-6            [2, 4]                    --
├─ModuleDict (block1): 1-3               --                        --       (block1_children=[Linear: 2, Linear: 2, Linear: 2]
│    └─Linear (in_lin4): 2-7             [2, 4]                    20
│    └─Linear (in_lin5): 2-8             [2, 4]                    20
│    └─Linear (in_lin6): 2-9             [2, 4]                    20
├─ModuleDict (block0): 1-4               --                        (recursive)       (block0_children=[Linear: 2]
│    └─Linear (in_lin1): 2-10            [2, 4]                    (recursive)
├─Linear (out_lin7): 1-5                 [2, 2]                    10
==========================================================================================
Total params: 646
Trainable params: 646
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
==========================================================================================