How to initialize layers for a Pytorch multi-task model?

username: vampire-tap
cooked: <p>I wish to have a multi-task model for two tasks in Pytorch. I coded it as follows:</p>
<pre><code class="lang-auto">x_trains=[]
y_trains=[]

num_tasks=2
for i in range(num_tasks):
    x_trains.append(torch.from_numpy(np.random.rand(100,1,50)).float())
    y_trains.append(torch.from_numpy(np.array([np.random.randint(10) for i in range(100)])).long())
nb_classes=10

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.all_task_layers=[]
        for i in range(num_tasks):
            self.all_task_layers.append(nn.Conv1d(1, 128, 8))
            self.all_task_layers.append(nn.BatchNorm1d(128))
            self.all_task_layers.append(nn.Conv1d(128, 256, 5))
            self.all_task_layers.append(nn.BatchNorm1d(256))
            self.all_task_layers.append(nn.Conv1d(256, 128, 3))
            self.all_task_layers.append(nn.BatchNorm1d(128))
            self.all_task_layers.append(nn.Linear(128, nb_classes))
        #self.dict_layers_for_tasks[i][1]        
        self.all_b1s=[]
        self.all_b2s=[]
        self.all_b3s=[]
        self.all_dense1s=[]
        
    def forward(self, x_num_tasks): 
        for i in range(0,len(self.all_task_layers),num_tasks):
            self.all_b1s.append(F.relu(self.all_task_layers[i+1](self.all_task_layers[i+0](x_num_tasks[i]))))

        for i in range(0,len(self.all_task_layers),num_tasks):
            self.all_b2s.append(F.relu(self.all_task_layers[i+3](self.all_task_layers[i+2](self.all_b1s[i]))))
                        
        for i in range(0,len(self.all_task_layers),num_tasks):
            self.all_b3s.append(F.relu(self.all_task_layers[i+5](self.all_task_layers[i+4](self.all_b2s[i]))))

        for i in range(0,len(self.all_task_layers),num_tasks):
            self.all_dense1s.append(self.all_task_layers[i+6](self.all_b3s[i]))
            
        return self.all_dense1s       

model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

losses=[]
for t in range(50):
    y_preds= model(x_trains)
    optimizer.zero_grad()
    for i in range(num_tasks):
        loss=criterion(y_preds[i], y_trains[i])
        losses.append(loss)
        print(t,i,loss.item())
        loss.backward(retain_graph=True)
        optimizer.step()
</code></pre>
<p>But it gives me error for initializing the model.parameters(). Specifically I get an empty list. The same model if I try to initialize the layers for the two tasks with two sets of Conv-&gt;BatchNorm-&gt;Conv-&gt;BatchNorm-&gt;Conv-&gt;BatchNorm-&gt;GlobalAveragePooling-&gt;Linear then it would work. But if I have let’s say 5 tasks, then it can get pretty messy, and that is why  I created a list of layers, and indexed them. For example, the first 7 layers are for task 1 and then the last 7 for task 2. But then the model.parameters() is giving me an empty list. How else could I do it? Or is there a simple fix, I am overlooking?</p>

username: ptrblck
cooked: <p>Try to use <code>nn.ModuleList</code> instead of Python <code>lists</code>, as this will properly register the modules internally.</p>

username: vampire-tap
cooked: <p>yup this worked <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> thanks alot!</p>