我有两个网络需要在我的完整模型中连接。然而,我的第一个模型是预训练的,在训练完整模型时需要将其设置为不可训练。我该如何在PyTorch中实现这一点。
我可以使用 this answer 来连接两个模型。
class MyModelA(nn.Module):
def __init__(self):
super(MyModelA, self).__init__()
self.fc1 = nn.Linear(10, 2)
def forward(self, x):
x = self.fc1(x)
return x
class MyModelB(nn.Module):
def __init__(self):
super(MyModelB, self).__init__()
self.fc1 = nn.Linear(20, 2)
def forward(self, x):
x = self.fc1(x)
return x
class MyEnsemble(nn.Module):
def __init__(self, modelA, modelB):
super(MyEnsemble, self).__init__()
self.modelA = modelA
self.modelB = modelB
def forward(self, x):
x1 = self.modelA(x)
x2 = self.modelB(x1)
return x2
# Create models and load state_dicts
modelA = MyModelA()
modelB = MyModelB()
# Load state dicts
modelA.load_state_dict(torch.load(PATH))
model = MyEnsemble(modelA, modelB)
x = torch.randn(1, 10)
output = model(x)
基本上,我想加载预训练的
modelA
并在训练Ensemble模型时使其无法训练。