PyTorch|保存及加载模型、nn.Sequential、ModuleList和ModuleDict

系列文章目录

PyTorch|Dataset与DataLoader使用、构建自定义数据集
PyTorch|搭建分类网络实例、nn.Module源码学习
pytorch|autograd使用、训练模型

文章目录


一、保存及加载模型

通过torch.save可以将该模型的参数、优化器状态、batch normalization、dropout、buffer变量等信息。

python 复制代码
import torch
import torchvision.models as models

(一)保存及加载模型的权重

模型取自torchvision.models里的vgg16,权重为IMAGENET1K_V1。

model.state_dict()是模型的权重。state_dict状态字典:一般包含当前model的参数及buffer变量

python 复制代码
model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')

推理时可以实现模型的加载:

  • 创建模型实例
  • 将实现保存的模型信息通过torch.load导入进来
  • 采用load_state_dict函数将模型信息载入模型实例
  • model.eval()使得模型进入推理模式
python 复制代码
model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

(二)保存及加载优化器的权重

保存优化器权重:

加载优化器权重:

(三)保存及加载整个模型

保存整个模型:

python 复制代码
torch.save(model, 'model.pth')

加载整个模型:

python 复制代码
model = torch.load('model.pth')

(四)保存及加载更具一般性的checkpoint

保存并加载用于推理或恢复训练的一般性checkpoint有助于从上次中断的地方重新开始。在保存一般检查点时,不仅仅是保存模型的state_dict,还包括保存优化器的state_dict、停止使用的时间,最近记录的训练损失,外部的torch.nn.Embedding层等等。

python 复制代码
# Additional information
EPOCH = 5
PATH = "model.pt"
LOSS = 0.4

torch.save({
            'epoch': EPOCH,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS,
            }, PATH)

加载:

python 复制代码
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

(五)保存多个模型

保存多个模型时可以将其直接合并到一个大字典中保存。

python 复制代码
# Specify a path to save to
PATH = "model.pt"

torch.save({
            'modelA_state_dict': netA.state_dict(),
            'modelB_state_dict': netB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            }, PATH)

二、nn.Sequential源码分析

nn.Sequential是有序的,当实例化nn.Sequential时,传入的模块顺序就是神经网络前向传播的顺序

在使用nn.Sequential时,可以按顺序传入模块,也可以输入一个字典。

(一)init函数

如果输入的是一个字典,init函数会采用遍历字典的方式,如果是一个一个的模块,init函数也会针对性的采取其他遍历方法。

(二)forward函数

对于一个模型的输入,nn.Sequential会依次的过其中的子模块。

nn.Sequential相比于ModuleList和ModuleDict来说,优势在于具有forward的功能。

三、ModuleList和ModuleDict

(一)ModuleList

pytorch允许我们把很多子模块放到一个列表中。ModuleList就是用于存放多个子模块的一个列表,在使用时可以对其进行遍历。ModuleList不单纯是一个列表,它本身就是一个module。

(二)ModuleDict

ModuleDict是用于存放多个子模块的一个字典,在使用时可以根据索引获得对应的子模块。ModuleDict不单纯是一个字典,它本身也是一个module。

除此之外,还有ParameterList、ParameterDict等,这些与ModuleList和ModuleDict的作用及使用方式类似。

参考:
8、深入剖析PyTorch的state_dict、parameters、modules源码
9、深入剖析PyTorch的nn.Sequential及ModuleList源码

相关推荐
冷眼看人间恩怨几秒前
【话题讨论】AI大模型重塑软件开发:定义、应用、优势与挑战
人工智能·ai编程·软件开发
2401_883041082 分钟前
新锐品牌电商代运营公司都有哪些?
大数据·人工智能
魔道不误砍柴功6 分钟前
Java 中如何巧妙应用 Function 让方法复用性更强
java·开发语言·python
_.Switch31 分钟前
高级Python自动化运维:容器安全与网络策略的深度解析
运维·网络·python·安全·自动化·devops
AI极客菌1 小时前
Controlnet作者新作IC-light V2:基于FLUX训练,支持处理风格化图像,细节远高于SD1.5。
人工智能·计算机视觉·ai作画·stable diffusion·aigc·flux·人工智能作画
阿_旭1 小时前
一文读懂| 自注意力与交叉注意力机制在计算机视觉中作用与基本原理
人工智能·深度学习·计算机视觉·cross-attention·self-attention
王哈哈^_^1 小时前
【数据集】【YOLO】【目标检测】交通事故识别数据集 8939 张,YOLO道路事故目标检测实战训练教程!
前端·人工智能·深度学习·yolo·目标检测·计算机视觉·pyqt
测开小菜鸟2 小时前
使用python向钉钉群聊发送消息
java·python·钉钉
Power20246662 小时前
NLP论文速读|LongReward:基于AI反馈来提升长上下文大语言模型
人工智能·深度学习·机器学习·自然语言处理·nlp
数据猎手小k2 小时前
AIDOVECL数据集:包含超过15000张AI生成的车辆图像数据集,目的解决旨在解决眼水平分类和定位问题。
人工智能·分类·数据挖掘