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

I wish to have a multi-task model for two tasks in Pytorch. I coded it as follows:

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()

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->BatchNorm->Conv->BatchNorm->Conv->BatchNorm->GlobalAveragePooling->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?

Try to use nn.ModuleList instead of Python lists, as this will properly register the modules internally.

yup this worked :slight_smile: thanks alot!