时间戳

首页 > 解决方案 > 向预训练模型添加参数

pytorch - 向预训练模型添加参数

问题描述

在 Pytorch 中，我们加载预训练模型如下：

net.load_state_dict(torch.load(path)['model_state_dict'])

那么网络结构和加载的模型必须完全相同。但是，是否可以加载权重然后修改网络/添加额外的参数？

注意：如果我们在加载权重之前向模型添加一个额外的参数，例如

self.parameter = Parameter(torch.ones(5),requires_grad=True)

Missing key(s) in state_dict:加载权重时会出错。

标签： pytorch

解决方案

让我们创建一个模型并保存它的状态。

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

        self.encoder = nn.LSTM(100, 50)

    def forward(self):
        pass


model1 = Model1()
torch.save(model1.state_dict(), 'filename.pt') # saving model

然后创建第二个模型，它与第一个模型有几个共同的层。加载第一个模型的状态并将其加载到第二个模型的公共层。

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

        self.encoder = nn.LSTM(100, 50)
        self.linear = nn.Linear(50, 200)

    def forward(self):
        pass


model1_dict = torch.load('filename.pt')
model2 = Model2()
model2_dict = model2.state_dict()

# 1. filter out unnecessary keys
filtered_dict = {k: v for k, v in model1_dict.items() if k in model2_dict}
# 2. overwrite entries in the existing state dict
model2_dict.update(filtered_dict)
# 3. load the new state dict
model2.load_state_dict(model2_dict)

推荐阅读