如何用另一个冻结模型训练pytorch模型?

Question

我用预训练模型(name = B)训练pytorch模型(name = A)。

我想像这样串联使用两个模型:使用输入，以便通过两个模型输出输出，并比较该输出和目标标签。Model2冻结并学习model1的学习参数。

图1 (b).

我这样写代码:

Class A(nn.Module):
# MY CODE
Class B(nn.Module):
# MY CODE
train_model = A()
freeze_model = B()
freeze_model.load_state_dict(torch.load("bestmodel.pth"))
for param in freeze_model.parameters():
param.requires_grad = False
criterion=nn.MSELoss()
optimizer = torch.optim.Adam(train_model.parameters(), lr=learning_rate)
for epoch in range(100):
... iteration ...
out1 = train_model(input1)
out2 = freeze_model(out1,input2)
loss = criterion(target,out2)
... optimizer.zero_grad, loss backward, ...
optimizer.step()

然而，第一次迭代得到了一个实值，在optimizer_Step之后，模型权重和损失被更改为nan..

我认为这是一个优化器或其他原因。有解决办法吗?

Answer 1

我得到了你的问题:
你不能在单个优化器上使用2个模型，使用这个

out1 = train_model(input1)
out2 = freeze_model(out1,input2)

我正在回答你的问题，但我只能写签名，你必须做必要的修改

Class A(nn.Module):
# MY CODE
Class B(nn.Module):
# MY CODE
train_model = A()
freeze_model = B()
freeze_model.load_state_dict(torch.load("bestmodel.pth"))
for param in freeze_model.parameters():
param.requires_grad = False

class Serial(nn.Module):
def __init__(self):
super(VGG, self).__init__()

# get the pretrained VGG19 network
self.train_1 = train_model
self.freeze_1 = freeze_model

def forward(self, x,input_2):
x = self.train_1(x)
x = self.freeze_1(x,input_2)
return x
serial_model = Serial()       
criterion=nn.MSELoss()
optimizer = torch.optim.Adam(serial_model.parameters(), lr=learning_rate)
for epoch in range(100):
... iteration ...
out = serial_model(input1,input2)
loss = criterion(target,out)
... optimizer.zero_grad, loss backward, ...
optimizer.step()