nn.Parameter is ommitted (with a case)

Describe the bug nn.Parameter is omitted in summary when there are other pytorch predefined layers in the networks. Details are as follows:

To Reproduce

import torch
import torch.nn as nn
from torchinfo import summary

class FCNets(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        # 2 layer fully connected networks
        super().__init__()
        # layer1 with nn.Parameter
        self.weight = nn.Parameter(torch.randn(input_dim, hidden_dim))
        self.bias = nn.Parameter(torch.randn(hidden_dim))
        # layer2 with nn.Linear
        self.fc2  = nn.Linear(hidden_dim, output_dim)
        # activation
        self.activation = nn.ReLU()
    
    def forward(self, x):
        # x.shape = [batch_size, input_dim]
        # layer1
        h = torch.mm(x, self.weight) + self.bias
        # activation
        h = self.activation(h)
        # layer2
        out = self.fc2(h)
        return out

# device = torch.device("cuda:0")
device = torch.device("cpu")
x = torch.randn(3, 128).to(device)
fc = FCNets(128, 64, 32).to(device)
summary(fc, input_data=x)

It seems that nn.Parameter is not compatible with other layers (nn.Module class).

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
FCNets                                   --                        --
├─ReLU: 1-1                              [3, 64]                   --
├─Linear: 1-2                            [3, 32]                   2,080
==========================================================================================
Total params: 2,080
Trainable params: 2,080
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.01
Estimated Total Size (MB): 0.01
==========================================================================================

However, if we remove self.fc2, the output will be fine.

Pytorch version: 1.7.1 (GPU) Torchinfo version: 1.5.3

[TylerYep]

This has been fixed in https://github.com/TylerYep/torchinfo/commit/5621cc9c316f2462b01a1392975212ac263f3fa1 and will be released in torchinfo v1.7.0.

Thank you for reporting this issue!

Also, thank you to @jeremyfix for the help on investigating. The above discussion was very helpful in diagnosing and fixing this issue as well as the 4+ other issues related to this problem. It required a fairly complex rewrite of the library and the code is now in a much better place.

[jeremyfix]

Ok, I will have to withdraw the work on this issue. I made some progress in the output but I’m not very happy with the code I’m producing as i’m having a hard time understanding the logic behind so I prefer to stop there.

My current progress in on the branch https://github.com/jeremyfix/torchinfo/tree/issue84 in which I basically added the computation of the number of parameters for the topmost level module (the one provided to the summary() function (code taken from the LayerInfo class). I also added two tests that are passed. I thought finally that adding a hook on the top level module was probably not necessary as we would call the forward on this function anyway and we just wanted to add the eventual parameters if any in the summary list. Also, when I was adding the hook, for some reasons, although the parameters got listed, the total number of parameters was not updated and I had to filter the named_parameters to just keep the top level ones.

In the current version of the code, 2 tests of the test suite are failing because the var_name of the topmost level module is displayed while it should be not according to the expected output

• Test test_row_settings because displaying LSTMNet (LSTMNet) instead of LSTMNet
• Test test_lstm because displaying SingleInputNet (SingleInputNet) instead of SingleInputNet

These are certainly minor things but probably required to dig into formatting.py to understand why.

The FCNets_ablin, FCNets_ab example above outputs almost the correct output :

-------- Result 1   
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
FCNets_ablin                             --                        --
├─a                                                                ├─8,192
├─b                                                                └─64
├─Linear: 1-1                            [3, 32]                   2,080
│    └─weight                                                      ├─2,048
│    └─bias                                                        └─32
==========================================================================================
Total params: 10,336
Trainable params: 10,336
Non-trainable params: 0
Total mult-adds (M): 0.01
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.04
Estimated Total Size (MB): 0.04
==========================================================================================

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
FCNets_ab                                [3, 64]                   8,256
├─a                                                                ├─8,192
├─b                                                                └─64
==========================================================================================
Total params: 8,256
Trainable params: 8,256
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.03
Estimated Total Size (MB): 0.04
==========================================================================================

“Almost correct” because I did not handle in the code the computation of the MACS, therefore the Total mult-adds and Estimated Total Size outputed are incorrect.

This can more easily be seen from the expected output of the parameter_with_other_layers.out where two exactly identical networks are summarized and the test does not succeed because these fields have incorrect values.

I’m sorry to resign but I hope that piece of work might be nevertheless of help for you.