Issue
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])
Solution
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.
Answered By - awaelchli
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.