Pytorch：如何训练具有两个损失函数的网络？

Question

我想先用重建损失预训练一个网络，然后通过交叉熵损失对其进行微调。 但似乎我必须在这两个阶段定义两个网络。 如何实现？

class Net():
    def __init__(self,pretrain):
        self.pretrain = pretrain
    def encoder(self,x):
        # do something here
        return x
    def decoder(self,x):
        # do something here
        return x
    
    def forward(self):
        e_x = self.encoder(x)
        if self.pretrain:
            return decoder(e_x)
        else:
            return e_x

def train(x,y):
    pretrain = True
    if pretrain:
        network = Net(pretrain=True)
        output = network(x)
        loss = MSE(x,output)
     else:
        network = Net(pretrain=False)
        output = network(x)
        loss = crossentropy(output,y)
    loss.backward()

Answer 1

您可以通过简单地定义两个损失函数来实现这一点，而 loss.backward 将适用于 go。 请参阅此处的相关讨论

MSE = torch.nn.MSELoss()
crossentropy = torch.nn.CrossEntropyLoss()
   
def train(x,y):
        pretrain = True
        if pretrain:
            network = Net(pretrain=True)
            output = network(x)
            loss = MSE(x,output)
        else:
            network = Net(pretrain=False)
            output = network(x)
            loss = crossentropy(output,y)
        loss.backward()