Pytorch加载部分预训练模型的参数

问题背景

假设我有一个已训练好的Model1，并已保存它的参数为.pth格式，我有一个与Model1结构完全相同的模型Model2，我希望Model2加载Model1中与特征提取有关的模块的参数，其他模块的参数随机初始化。

应用场景为在K折交叉验证时，我希望从第二折开始的模型加载第一折训练模型的部分参数，并在此基础上微调，从而减少训练轮数。

解决方法

加载保存的第一折训练好的模型参数，因为我保存时是多GPU，加载时需要多GPU模型参数去掉module：

state_dict = torch.load("./Fold_1_checkpoint_4.pth",map_location=torch.device('cpu'))
# 多GPU 模型参数去掉 module
from collections import OrderedDict
state_dict_new = OrderedDict()
for k, v in state_dict.items():
    name = k[7:]  # 去掉 `module.`
    # name = k.replace(“module.", "")
    state_dict_new[name] = v

打印出来看看：

OrderedDict([('LD1.conv.weight',
              tensor([[[ 0.0047,  0.0013,  0.0026,  0.0025,  0.0060, -0.0055,  0.0024,
                        -0.0022,  0.0052, -0.0039, -0.0015,  0.0049,  0.0861,  0.0095,
                         0.0032,  0.0033,  0.0037,  0.0008,  0.0023,  0.0080,  0.0045,
                        -0.0023,  0.0097,  0.0034,  0.0160]]])),
             ('LD1.conv.bias', tensor([0.0101])),
             ('position_embedder1.position_embedding',
              tensor([[-0.1612, -0.2091, -0.9842,  ...,  0.7624, -0.7286, -0.6549],
                      [-0.2659,  0.0225, -1.1011,  ...,  1.3899, -0.6597, -0.7806],
                      [-0.3426, -0.2478, -1.0329,  ...,  1.2922, -0.8009, -0.8948],
                      ...,
                      [ 1.4002,  0.0854,  0.4595,  ...,  0.8604, -0.6667, -0.3106],
                      [-0.2322,  0.1524,  0.0809,  ..., -0.3961,  1.2030, -0.3428],
                      [-0.7272,  2.0944, -0.7098,  ...,  0.2923,  0.1994, -0.2035]])),
             ······
             ('encoder1.attn_layers.0.attention.gen1.weight',
              tensor([[-0.2019,  0.2232,  0.2444,  ..., -0.0791,  0.1458,  0.2732],
                      [-0.3740,  0.0378, -0.1596,  ..., -0.2621,  0.2512,  0.1971],
                      [ 0.0941, -0.2160, -0.1651,  ...,  0.1103,  0.2440, -0.8079],
                      ...,
                      [-0.0320, -0.2835,  0.7415,  ...,  0.1435, -0.5525, -0.2343],
                      [-0.2918, -0.4909,  0.0787,  ..., -0.2581, -0.2801,  0.5388],
                      [-0.3744, -0.8789, -0.3549,  ..., -0.0468, -0.1703, -0.1652]])),
             ('encoder1.attn_layers.0.attention.gen1.bias',
              tensor([ 0.1687, -0.8035, -0.0674, -0.4795, -0.2964, -0.5702, -0.9225, -0.0856,
                       1.4210, -0.9099, -1.3605, -0.8256, -1.6564, -1.5394, -0.7184,  0.4084,
                      -0.6281, -2.2652, -0.7902, -0.9841, -0.4321, -0.4909, -1.1394, -1.2694,
                      -0.0219, -0.3581,  0.5682, -1.2015, -0.9408, -1.4014, -1.2602, -1.3793,
                      -0.8697, -2.1542, -0.7590, -1.7175,  0.0472,  0.0683,  0.4771, -1.7388,
                      -1.7942, -0.2653, -0.2880,  0.1545, -1.4922, -0.3756,  0.5862, -1.5961,
                      -0.2772,  0.9524, -0.1537, -0.4896, -0.2866, -0.9232,  0.2729,  0.0779,
                      -1.7396,  0.0744, -0.9369,  0.3411, -0.2821, -0.9431, -1.1710,  0.1081])),
              ······
             ('gate1.0.weight',
              tensor([[ 0.4562, -0.0911,  0.2505,  ...,  0.4371,  0.1663, -0.4728],
                      [ 0.1472,  0.1287,  0.0289,  ...,  0.0881,  0.0862, -0.2207],
                      [ 0.5243,  0.0716,  0.4541,  ..., -0.4321,  0.6488,  0.7909],
                      ...,
                      [-0.0867,  0.2686,  0.1457,  ...,  0.0756,  0.4644, -0.4607],
                      [-0.1575,  0.3019,  0.1621,  ..., -0.1950,  0.2406, -0.2090],
                      [-0.0231, -0.2043, -0.0222,  ...,  0.1126,  0.2387, -0.7590]])),
             ('gate1.0.bias',
              tensor([-0.5651, -0.6304, -0.6943, -0.5891, -0.2952, -0.3897, -0.6592, -0.3604,
                       0.1627, -0.4173, -0.3456, -0.6569, -0.0337, -0.1433, -0.3954, -0.0141])),
             ······
             ('Linear_res2.bias',
              tensor([-0.0314, -0.0065,  0.0268, -0.0321, -0.0738, -0.0063, -0.0531, -0.0615,
                      -0.0552, -0.0357, -0.0639, -0.0893, -0.0361, -0.0736, -0.0347, -0.0330,
                      -0.0759, -0.0828, -0.0665, -0.0439, -0.0652, -0.0718, -0.0231, -0.0297,
                      -0.0448, -0.0408, -0.0181, -0.0379, -0.0274, -0.0526, -0.0139, -0.0404,
                      -0.0284, -0.0496, -0.0515, -0.0054, -0.0704, -0.0666, -0.0385, -0.0613,
                      -0.0471, -0.0886, -0.0398, -0.0616, -0.0304, -0.0558, -0.0301, -0.0728,
                      -0.0869, -0.0409, -0.0514, -0.0737, -0.0510, -0.1048, -0.0555, -0.0530,
                      -0.0721, -0.0315, -0.0070, -0.0687, -0.0707, -0.0403, -0.0611, -0.0340,
                      -0.0935, -0.0339, -0.0462, -0.0842, -0.0516, -0.0445, -0.0364, -0.0748]))])

在此案例中，模型特征提取部分为encoder开头和gate开头的模块，因此需要过滤掉其他不需要迁移的模块：

pretrained_dict = {k: v for k, v in state_dict_new.items() if k.split('.')[0][:-1] in ['encoder','gate']}

初始化Model2：

model2 = Model(args).cuda()
model_dict = model2.state_dict()

更新model_dict中需要迁移的部分并导回模型：

model_dict.update(pretrained_dict)
model2.load_state_dict(model_dict)

至此，model2中就加载了model1中与特征提取相关的模块参数，即可在此基础上微调。

参考文献

Pytorch如何加载部分预训练模型的参数
pytorch如何使模型只更新一部分参数 pytorch加载模型部分参数 转载
pytorch加载多GPU模型和单GPU模型