Pytorch加载预训练权重的方法

Pytorch加载预训练权重

使用Pytorch加载预训练权重时，可能会出现以下两个问题：

1. 预训练网络和目标网络的结构不同，权重形状不一致
2. 预训练网络和目标网络的网络层名称不同，结构相同

直接加载时以上问题都会导致报错，下面的代码可以实现正确加载，只加载网络层名称和权重形状一致的层，不一致时保留目标网络的。

import torch
from typing import OrderedDict


def load_pretrained(dict1: OrderedDict, dict2: OrderedDict):
    """
    加载预训练权重，键（网络层名）或值（网络权重）的形状不一致时保留dict2的键值
    :param dict1: 预训练网络的权重
    :param dict2: 目标网络的权重
    :return: dict2
    """
    assert isinstance(dict1, OrderedDict), f"dict1's type is {type(dict1)},not OrderedDict"
    assert isinstance(dict2, OrderedDict), f"dict2's type is {type(dict2)},not OrderedDict"
    for k, v in dict1.items():
        if k in dict2 and v.shape == dict2[k]:
            dict2[k] = v
    return dict2


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
prenet = torch.load('预训练权重文件的路径*.pt')
prenet_state_dict = prenet.state_dict() # 该行不固定，不同的权重文件保存的键值对不同，要选出模型的权重
net = Net() # 目标网络
net_dict = load_pretrained(prenet_state_dict , net.state_dict())
net.load_state_dict(net_dict, strict=False)
net.to(device)