CODEWITHSUNDEEP
pythonpytorch
vampiretap
vampiretap

Reputation: 361

How to define several layers via a loop in __init__ for Pytorch?

I am trying to define a multi-task model in Pytorch where I need a different set of layers for different tasks. I face problems in defining layers, especially if I use a for loop to store different layers in a list then I get an error from optimizer stating that model.parameters() is an empty list, which in fact it is.

Following is the code: 

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

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?

Upvotes: 6

Views: 8017

Answers (1)

cookiemonster
cookiemonster

Reputation: 2104

you should use nn.ModuleList() to wrap the list. for example
x_trains = nn.ModuleList(x_trains)

see PyTorch : How to properly create a list of nn.Linear()

Upvotes: 8

Related Questions