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

news2024/12/29 12:59:02

系列文章目录

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

文章目录

系列文章目录
一、保存及加载模型
- （一）保存及加载模型的权重
- （二）保存及加载优化器的权重
- （三）保存及加载整个模型
- （四）保存及加载更具一般性的checkpoint
- （五）保存多个模型
二、nn.Sequential源码分析
- （一）init函数
- （二）forward函数
三、ModuleList和ModuleDict
- （一）ModuleList
- （二）ModuleDict

一、保存及加载模型

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

import torch
import torchvision.models as models

（一）保存及加载模型的权重

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

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

model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')

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

创建模型实例
将实现保存的模型信息通过torch.load导入进来
采用load_state_dict函数将模型信息载入模型实例
model.eval()使得模型进入推理模式

model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

（二）保存及加载优化器的权重

保存优化器权重：
在这里插入图片描述

加载优化器权重：
在这里插入图片描述

（三）保存及加载整个模型

保存整个模型：

torch.save(model, 'model.pth')

加载整个模型：

model = torch.load('model.pth')

（四）保存及加载更具一般性的checkpoint

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

# 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)

加载：

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()

（五）保存多个模型

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

# 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)