Pytorch lightning: see input/ouptut size in model summary when using nn.ModuleList

ThomaS 0
Viewed 34

When I use nn.ModuleList() to define layers in my pytorch lightning model, their "In sizes" and "Out sizes" in ModelSummary are "?" Is their a way to have input/output sizes of layer in model summary, eventually using something else than nn.ModuleList() to define layers from a list of arguments.

Here is a dummy model: (12 is the batch size)

import torch
import torch.nn as nn
from pytorch_lightning import LightningModule
from pytorch_lightning.utilities.model_summary import ModelSummary


class module_list_dummy(LightningModule):
    def __init__(self,
                 layers_size_list,                 
                ):
        super().__init__()
        self.example_input_array = torch.zeros((12,100), dtype=torch.float32)
        self.fc11 = nn.Linear(100,50)
        self.moduleList = nn.ModuleList()
        input_size = 50
        for layer_size in layers_size_list:
            self.moduleList.append(nn.Linear(input_size, layer_size))
            input_size = layer_size
        self.loss_fn = nn.MSELoss()
        
    def forward(self, x):
        out = self.fc11(x)
        for layer in self.moduleList:
            out = layer(out)
        return out
    
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        out = self.fc11(X)
        for layer in self.moduleList:
            out = layer(out)
        loss = torch.sqrt(self.loss_fn(out, y))
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        out = self.fc11(X)
        for layer in self.moduleList:
            out = layer(out)
        loss = torch.sqrt(self.loss_fn(y_hat, y))
        self.log('val_loss', loss)
        return loss
   
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
        return optimizer

the following code print the summary:

net = module_list_dummy(layers_size_list=[30,20,1])
summary = ModelSummary(net)
print(summary)

But the output is:

| Name       | Type       | Params | In sizes  | Out sizes
------------------------------------------------------------------
0 | fc11       | Linear     | 5.0 K  | [12, 100] | [12, 50] 
1 | moduleList | ModuleList | 2.2 K  | ?         | ?        
2 | loss_fn    | MSELoss    | 0      | ?         | ?        
------------------------------------------------------------------
7.2 K     Trainable params
0         Non-trainable params
7.2 K     Total params
0.029     Total estimated model params size (MB)

I would expect to have:

1 | moduleList | ModuleList | 2.2 K  | [12,50]   | [12,1]

or even better something like

1  | Linear | Linear | ... | [12,50] |  [12,30]
2  | Linear | Linear | ... | [12,30] |  [12,20]
3  | Linear | Linear | ... | [12,20] |  [12,1]

I did this to check that those layer are being used in forward:

x = torch.zeros((12,100), dtype=torch.float32)
net(x).shape

and they are ( model output is size [12,1])

pytorch-lightning pytorch
1 Answers

In your case you are using the layers in the list in a sequential manner:

for layer in self.moduleList:
    out = layer(out)

However, nn.ModuleList does not force you to run your layers sequentially. You could be doing this:

out = self.moduleList[3](out)
out = self.moduleList[1](out)
out = self.moduleList[0](out)

Hence, ModuleList cannot itself be interpreted as a "layer" with an input- and output shape. It is completely arbitrary how the user might integrate the layers inside the list in their graph.

In your specific case where you run the layers sequentially, you could define a nn.Sequential instead. This would look like this:

self.layers = nn.Sequential()
for layer_size in layers_size_list:
    self.layers.append(nn.Linear(input_size, layer_size))
    input_size = layer_size

And then in your forward:

out = self.layers(input)

Now, the model summary in Lightning will also show the shape sizes, because nn.Sequential's output shape is now directly the output shape of the last layer inside of it.

I hope this answer helps your understanding.

awaelchli0
Related