我想先用重建损失预训练一个网络,然后通过交叉熵损失对其进行微调。但看来我必须在这两个阶段定义两个网络。如何实现?
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()
您可以通过简单地定义两个损失函数来实现这一点,并且 loss.backward 就可以了。请参阅相关讨论这里
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()