我想在两个相关的不同输入和目标的数据集上构建一个多任务学习模型。这两个任务共享底层层,但有不同的头部层,一个最小的示例:
class MultiMLP(nn.Module): """ A simple dense network for MTL on hard parameter sharing. """ def __init__(self): super().__init__() self.hidden = nn.Linear(100, 200) self.out_task0= nn.Linear(200, 1) self.out_task0= nn.Linear(200, 1) def forward(self, x): x = self.hidden(x) x = F.relu(x) y_task0 = self.out_task0(x) y_task1 = self.out_task1(x) return [y_task0, y_task1]
数据加载器的构造方式使得批次交替地从两个数据集中生成,即批次0, 2, 4, …来自任务0,批次1, 3, 5, …来自任务1。我想以这种方式训练网络:仅针对来自任务0的批次更新hidden
层和out_task0
的权重,而仅针对任务1更新hidden
和out_task1
的权重。
然后我在训练过程中交替切换requires_grad
,如下所示。但我观察到每次迭代时所有权重都会被更新。
...criterion = MSELoss()for i, data in enumerate(combined_loader): x, y = data[0], data[1] optimizer.zero_grad() # controller is 0 for task0, 1 for task1 # altenate the header layer controller = i % 2 task0_mode = True if controller == 0 else False for name, param in model.named_parameters(): if name in ['out_task0.weight', 'out_task0.bias']: param.requires_grad = task0_mode elif name in ['out_task1.weight', 'out_task1.bias']: param.requires_grad = not task0_mode outputs = model(x)[controller] loss = criterion(outputs, y) loss.backward() optimizer.step() # Monitor the parameter updates for name, p in model.named_parameters(): if name in ['out_task0.weight', 'out_task1.weight']: print(f"Controller: {controller}") print(name, p)
我在训练过程中遗漏了什么吗?还是整个设置不会工作?
回答:
免责声明:这个问题已在PyTorch论坛上得到解答,我在这里整理内容以防其他人遇到相同的问题,感谢ptrblk
问题可能源于任何随机梯度下降(sgd)的变体,这些变体利用了之前步骤的梯度,例如带动量的随机梯度下降(sgd-m)、Nesterov加速梯度(NAG)、Adagrad、RMSprop、Adam等。在步骤t
清零梯度不会影响依赖历史梯度的项。因此,在提出的问题中的设置下,权重仍然会被更新。
可以从以下代码示例中看到这一点。
model = nn.Linear(1, 1, bias=False)#optimizer = torch.optim.SGD(model.parameters(), lr=1., momentum=0.) # same results for w1 and w2optimizer = torch.optim.SGD(model.parameters(), lr=1., momentum=0.5) # w2 gets updated#optimizer = torch.optim.Adam(model.parameters(), lr=1.) # w2 gets updatedw0 = model.weight.clone()out = model(torch.randn(1, 1))out.mean().backward()optimizer.step()w1 = model.weight.clone()optimizer.zero_grad()print(model.weight.grad)optimizer.step()w2 = model.weight.clone()print(w1 - w0)print(w2 - w1)
使用原生SGD优化器,w2和w1是相同的。但对于SGD-M和Adam则不然。