ModuleList not showing up in summary

Hi, it looks like modules wrapped in a ModuleList don’t show up in the summary (looks like the params are counted in the total though).

Minimal example:

import torch
from torchsummary import summary


class TestModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self._layers = torch.nn.ModuleList()
        self._layers.append(torch.nn.Linear(5,5))
        self._layers.append(torch.nn.Linear(5,5))

    def forward(self, x):
        out = x
        for l in self._layers:
            out = l(x)
        return out


if __name__ == "__main__":
    my_module = TestModule()
    summary(my_module, (5,))

Gives

------------------------------------------------------------------------------------------
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
==========================================================================================
Total params: 60
Trainable params: 60
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
------------------------------------------------------------------------------------------

Whereas the print(str(model)) gives

TestModule(
  (_layers): ModuleList(
    (0): Linear(in_features=5, out_features=5, bias=True)
    (1): Linear(in_features=5, out_features=5, bias=True)
  )
)

Is this expected behavior? I could probably work on a PR if you think this should be changed and you have no quick fix for it.

[TylerYep]

This should be changed! If you would like to start a PR, that would be awesome!

[TylerYep]

The changes look really great! I left a few minor comments on the PR. Functionality-wise it looks good to merge, but it would be great if you could add tests for all of the models you tried.

To do so, just add models to fixtures/models.py, and add each call to torchsummary_test.py. Output tests are optional, since the ModuleList output is subject to change.

I will think about the indexing strategy a bit more, but seeing as none of the previous test cases have broken, these changes seem good to merge - I’ll add a note that Sequential and ModuleList support is in-progress.