Pytorch预训练模型和修改——记录

最新推荐文章于 2024-07-02 21:03:46 发布

原创

最新推荐文章于 2024-07-02 21:03:46 发布 · 1.3k 阅读

12 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #机器学习 #模型参数修改

本文介绍如何在PyTorch中加载预训练模型及自定义模型，详细讲解了模型参数的修改方法，特别是针对全连接层的调整，并探讨了如何通过冻结特定层来进行微调。

文章目录

加载模型
读取预训练模型
模型参数修改
训练特定层（冻结层）
PyTorch的Module.modules()和Module.children()

加载模型

一般从torchvision的models中加载常用模型，如alexnet、densenet、inception、resnet、squeezenet、vgg等常用网络结构，并提供预训练模型，调用方便。

from torchvision import models

resnet = models.resnet50(pretrain=True)
print(resnet)  # 打印网络结构

读取预训练模型

另一种是读取自己预训练模型，而不是使用官方自带。

import torch
resnet18 = models.resnet18(pretrained=False) #pretrained参数默认是False,为了代码清晰，最好还是加上参数赋值.

resnet18.load_state_dict(torch.load(path_params.pkl))

load_state_dict方法还有一个重要的参数是strict，该参数默认是True，表示预训练模型的层和自己定义的网络结构层严格对应相等（比如层名和维度）。

当新定义的网络（model_dict）和预训练网络（pretrained_dict）的层名不严格相等时，需要先将pretrained_dict里不属于model_dict的键剔除掉：

pretrained_dict = {
   
   k: v for k, v in pretrained_dict.items() if k in model_dict}

再用预训练模型参数更新model_dict，最后用load_state_dict方法初始化自己定义的新网络结构。

整体代码：

print resnet18 #打印的还是网络结构
 
# 注意: cnn = resnet18.load_state_dict(torch.load( path_params.pkl )) #是错误的，这样cnn将是nonetype
 
pre_dict = resnet18.state_dict() #按键值对将模型参数加载到pre_dict