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

时间:2020-08-13 23:45:48

标签: python neural-network pytorch pre-trained-model

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

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

1 个答案:

答案 0 :(得分:1)

您可以通过简单地定义两个损失函数和损失来实现。参见相关讨论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()