Lnton羚通关于PyTorch的保存和加载模型基础知识

SAVE AND LOAD THE MODEL (保存和加载模型)

PyTorch 模型存储学习到的参数在内部状态字典中,称为 state_dict, 他们的持久化通过 torch.save 方法。

复制代码
model = models.shufflenet_v2_x0_5(pretrained=True)
torch.save(model, "../../data/ShuffleNetV2_X0.5.pth")

如果要加载模型的话,首先需要实例化一个同类型的模型对象,然后用 load_state_dict() 方法加载参数。

复制代码
model = models.shufflenet_v2_x0_5()
model.load_state_dict(torch.load("../../data/ShuffleNetV2_X0.5.pth"))
model.eval()

Output exceeds the size limit. Open the full output data in a text editor
ShuffleNetV2(
  (conv1): Sequential(
    (0): Conv2d(3, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (stage2): Sequential(
    (0): InvertedResidual(
      (branch1): Sequential(
        (0): Conv2d(24, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=24, bias=False)
        (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): Conv2d(24, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (3): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (4): ReLU(inplace=True)
      )
      (branch2): Sequential(
        (0): Conv2d(24, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): ReLU(inplace=True)
        (3): Conv2d(24, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=24, bias=False)
        (4): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (5): Conv2d(24, 24, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (6): BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (7): ReLU(inplace=True)
...
    (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace=True)
  )
  (fc): Linear(in_features=1024, out_features=1000, bias=True)
)

Saving and Loading Models with Shapes

当加载模型权重时,我们需要首先实例化模型类,因为类定义了网络的结构。我们可能想要保存类的结构以及模型,在这种情况下,我们可以将 model (而不是 model.state_dict() ) 传递给保存函数:

复制代码
torch.save(model, "../../data/ShuffleNetV2_X0.5_eval2.pth")

加载模型如这样:

复制代码
model = torch.load("../../data/ShuffleNetV2_X0.5_eval2.pth")
print(model)

这种方法在序列化模型时使用 Python pickle 模块,因此它依赖于加载模型时可用的实际类定义。

Lnton羚通专注于音视频算法、算力、云平台的高科技人工智能企业。 公司基于视频分析技术、视频智能传输技术、远程监测技术以及智能语音融合技术等, 拥有多款可支持ONVIF、RTSP、GB/T28181等多协议、多路数的音视频智能分析服务器/云平台。

相关推荐
Albert_Lsk1 分钟前
【2025/07/10】GitHub 今日热门项目
人工智能·开源·github·开源协议
巨人张6 分钟前
信息素养Python编程题
开发语言·python
新加坡内哥谈技术9 分钟前
OpenAI正准备推出一款搭载人工智能功能的网络浏览器,试图直接挑战Alphabet旗下
人工智能
站大爷IP11 分钟前
Python爬虫动态IP代理报错全解析:从问题定位到实战优化
python
hie9889418 分钟前
CentOS环境搭建-快速升级G++版本
linux·python·centos
盛寒21 分钟前
向量空间 线性代数
python·线性代数·机器学习
Better Rose22 分钟前
人工智能与机器学习暑期科研项目招募(可发表论文)
人工智能·深度学习·机器学习·论文撰写
Hoshino _Ai29 分钟前
OpenCV图像认知(三)
人工智能·opencv·计算机视觉
哈__30 分钟前
文心一言4.5开源部署指南及文学领域测评
人工智能·ai·文心一言·paddlepaddle
Jamence1 小时前
多模态大语言模型arxiv论文略读(155)
论文阅读·人工智能·计算机视觉·语言模型·论文笔记