Layer names

[apivovarov]

It would be useful to show Layer names in the grid. For example, print(model) shows me the following modules/Layers with names:

    (fpn_lateral2): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral3): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral4): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (fpn_lateral5): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
    (fpn_output5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (top_block): LastLevelMaxPool()

At the same time .summary shows:

├─Conv2d: 1-2                                 [1, 2048, 25, 38]         [1, 256, 25, 38]          [2048, 256, 1, 1]
├─Conv2d: 1-3                                 [1, 256, 25, 38]          [1, 256, 25, 38]          [256, 256, 3, 3]
├─Conv2d: 1-4                                 [1, 1024, 50, 76]         [1, 256, 50, 76]          [1024, 256, 1, 1]
├─Conv2d: 1-5                                 [1, 256, 50, 76]          [1, 256, 50, 76]          [256, 256, 3, 3]
├─Conv2d: 1-6                                 [1, 512, 100, 152]        [1, 256, 100, 152]        [512, 256, 1, 1]
├─Conv2d: 1-7                                 [1, 256, 100, 152]        [1, 256, 100, 152]        [256, 256, 3, 3]
├─Conv2d: 1-8                                 [1, 256, 200, 304]        [1, 256, 200, 304]        [256, 256, 1, 1]
├─Conv2d: 1-9                                 [1, 256, 200, 304]        [1, 256, 200, 304]        [256, 256, 3, 3]
├─LastLevelMaxPool: 1-10                      [1, 256, 25, 38]          [1, 256, 13, 19]          --

As you can see the order of Layers is different - it would be nice if we can identify Layers in the .summary output by name

[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"])