Pytorch: How to train a network with two loss functions?
I want to pretrain a network with reconstruction loss first, then finetune it by crossentropy loss. But it seems that I have to define two network in this two stage. How to achieve it?
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()


You can achieve this by simply defining the two-loss functions and loss.backward will be good to go. See the relevant discussion 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()

