pytorch权值更新中的交替训练多任务学习模型

Question

我想在两个具有不同输入和目标的相关数据集上构建一个多任务学习模型。这两个任务共享较低级别的层，但具有不同的头层，这是一个简单的示例：

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的eem>。我想以这种方式训练网络：只更新来自任务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)

我在培训过程中有没有遗漏什么？或者整个设置不起作用

Answer 1

免责声明：这个问题已经从PyTorch Forum得到了回答，我在这里把问题放在一起，以防有人遇到同样的问题，这要归功于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 w2
optimizer = torch.optim.SGD(model.parameters(), lr=1., momentum=0.5) # w2 gets updated
#optimizer = torch.optim.Adam(model.parameters(), lr=1.) # w2 gets updated

w0 = 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的情况并非如此