[Feature Request] Show children names as well

summary should have the option (default?) to show module/layer names as well (not just class name). I.e. use named_children when traversing module children. e.g. in the following example should print something like:

‘block1’ [Sequential]: ‘conv1 [Conv2d]’ ‘relu1 [ReLU]’ ‘block2’ [Sequential]: ‘conv1 [Conv2d]’ ‘relu1 [ReLE]’

from collections import OrderedDict
import torch
from torch import nn
import torchinfo

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.block1 = nn.Sequential(OrderedDict([
            ('conv1', nn.Conv2d(1, 8, 3, padding=1)),
            ('relu1', nn.ReLU())
            ]))
        self.block2 = nn.Sequential(OrderedDict([
            ('conv1', nn.Conv2d(8, 8, 3, padding=1)),
            ('relu1', nn.ReLU())
        ]))

    def forward(self, inputs):
        x = inputs['image']

        x = self.block1(x)
        x = self.block2(x)
        out = x.mean()

        return out


if __name__ == "__main__":
    model = Model()
    summary = torchinfo.summary(model)
    print(summary)

[m1pro]

For me it seems obvious. Different modules in the network have different functionalities. It makes it much more readable when the name of the sub module appears instead of just generic class name or Sequential. For example in computer vision models you can have feature extraction module, different heads of the same structure but different functionalities. For example Segmentation, classification, regressions. When one names a variable, one wants to see the name appear in a debugging windows not for example the address, right? 😉 And of course as the network grows bigger it makes it even harder when those labels are missing from the summary print.

[TylerYep]

I just added this today in version 0.1.0! You can add the layer names via:

torchinfo.summary(..., row_settings=["var_names", "depth"])