【学习1】Pytorch部分载入参数的几种方法

参考1 https://blog.csdn.net/LXX516/article/details/80124768

参考2 https://www.bilibili.com/video/BV1dA411g7Ao

方法一 仅仅修改了全连接层，比如将分类层fc的1000类输出改为10类输出的new_fc

import torch
from torchvision.models import resnet18
import torch.nn as nn

def main():
    out_features = 10 #输出的类别数
    net = resnet18(pretrained=True)
    in_features = net.fc.in_features
    #net.fc = nn.Linear(in_features,out_features)
    print(net)

方法二 对模型进行了部分修改，仅加载没有修改的部分的参数（修改的部分要用不同的名字，比如修改了fc层，修改后的层还叫fc用这个方法就会报错）

    pretrained_dict = torch.load(model_weight)  
    model_dict = myNet.state_dict()
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 
    model_dict.update(pretrained_dict)
    myNet.load_state_dict(model_dict)

方法三 最具有通用性的方法，弄清楚那些层中预训练的参数不要，直接不载入

    model_path = './resnet18.pth' #根据实际情况载入自己的模型参数
    net = resnet18(num_classes=5) #先设定了out_features为5，此时直接载入参数一定会报错
    pre_weights = torch.load(model_path,map_location='cpu') #先读取，不载入
    delete_key = []
    for key,_ in pre_weights.items():
        if 'fc' in key: #根据具体不需要的key值修改
            delete_key.append(key)

    for key in delete_key:
        del pre_weights[key]

    net.load_state_dict(pre_weights,strict=False) #strict默认是True，会严格按照key值来载入参数

方法四 实在有需要可以通过自己创建字典，在里面放入自己需要的key和value再载入到模型中（我自己也没尝试过，所以无代码）