Pytorch：如何训练具有两种损失功能的网络？

问题描述

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

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

解决方法

您可以通过简单地定义两个损失函数和损失来实现。参见相关讨论here

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