How can I access layers in a pytorch module by index?

Question

I am trying to write a pytorch module with multiple layers. Since I need the intermediate outputs I cannot put them all in a Sequantial as usual. On the other hand, since there are many layers, what I have in mind is to put the layers in a list and access them by index in a loop. Below describe what I am trying to achieve:

import torch
import torch.nn as nn
import torch.optim as optim

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.layer_list = []

        self.layer_list.append(nn.Linear(2,3))
        self.layer_list.append(nn.Linear(3,4))
        self.layer_list.append(nn.Linear(4,5))

    def forward(self, x):
        res_list = [x]
        for i in range(len(self.layer_list)):
            res_list.append(self.layer_list[i](res_list[-1]))
        return res_list


model = MyModel()
x = torch.randn(4,2)
y = model(x)

print(y)

optimizer = optim.Adam(model.parameters())

The forward method works fine, but when I want to set an optimizer the program says

ValueError: optimizer got an empty parameter list

It appears that the layers in the list are not registered here. What can I do?

Answer 1

If you put your layers in a python list, pytorch does not register them correctly. You have to do so using ModuleList ( https://pytorch.org/docs/master/generated/torch.nn.ModuleList.html ).

ModuleList can be indexed like a regular Python list, but modules it contains are properly registered, and will be visible by all Module methods.

Your code should be something like:

import torch
import torch.nn as nn
import torch.optim as optim

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

        self.layer_list = nn.ModuleList()  # << the only changed line! <<

        self.layer_list.append(nn.Linear(2,3))
        self.layer_list.append(nn.Linear(3,4))
        self.layer_list.append(nn.Linear(4,5))

    def forward(self, x):
        res_list = [x]
        for i in range(len(self.layer_list)):
            res_list.append(self.layer_list[i](res_list[-1]))
        return res_list

By using ModuleList you make sure all layers are registered in the computational graph.

There is also a ModuleDict that you can use if you want to index your layers by name. You can check pytorch's containers here: https://pytorch.org/docs/master/nn.html#containers

Answer 2

It is possible to list all layers on neural network by use

list_layers = model.named_children()

In the first case, you can use:

parameters = list(Model1.parameters())+ list(Model2.parameters())
optimizer = optim.Adam(parameters, lr=1e-3)

In the second case, you didn't create the object, so basically you can try this:

model = VAE()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

By the way, you can start from modifying the VAE example provided by Pytorch.

Perhaps you miss the initial function or initialize the model in a wrong way. See theinit function here.