0

0

Python中PyTorch如何保存和加载模型_使用state_dict序列化权重参数

P粉602998670

P粉602998670

发布时间:2026-03-19 10:37:45

|

454人浏览过

|

来源于php中文网

原创

应仅保存模型权重state_dict而非整个模型对象,因其不依赖类定义、支持跨环境加载,且需配合eval()模式、正确后缀、map_location及独立管理优化器状态。

python中pytorch如何保存和加载模型_使用state_dict序列化权重参数

保存模型权重用 torch.save(model.state_dict(), path),别存整个模型对象

直接保存 model 本身看似方便,但会把类定义、模块结构甚至训练时的临时变量一并固化,导致加载时对代码结构极度敏感——换一个文件路径、改一个 import 名、甚至只是升级了 PyTorch 版本,就可能报 AttributeError: 'dict' object has no attribute 'forward' 或更模糊的 Missing key 错误。

真正稳定、可迁移的做法是只序列化权重:用 model.state_dict() 提取参数字典,再用 torch.save() 写入磁盘。它不依赖模型类的具体实现,只要加载时能重建出结构一致的模型实例,就能安全 load_state_dict()

  • 保存前确保模型在 eval() 模式(避免 BatchNormDropout 的训练态参数被意外保存)
  • 路径建议用 .pt.pth 后缀,不要用 .pkl ——虽然底层都是 pickle,但后缀影响工具链识别(如 Hugging Face from_pretrained 默认找 .bin.pth
  • 如果模型用了 nn.DataParallel,保存前记得用 model.module.state_dict(),否则键名会多出 module. 前缀,后续单卡加载会不匹配

加载权重必须先初始化模型,再调用 load_state_dict()

不能跳过模型构造直接“反序列化”权重。PyTorch 不提供从 state_dict 自动还原网络结构的能力 —— 它只负责把字典里的值填进已有模型的对应参数位置。

常见错误是写成 model = torch.load(path),这会尝试 unpickle 整个对象,一旦环境不一致就崩溃;或者漏掉 model.load_state_dict(...),结果模型还是随机初始化的权重。

立即学习Python免费学习笔记(深入)”;

  • 加载后务必调用 model.eval()(如果用于推理),否则 BatchNorm 仍按训练模式运行,输出不稳定
  • 键名不匹配时默认报错,可用 strict=False 忽略多余或缺失的键,但要小心:missing keys 可能意味着模型结构没对齐,unexpected keys 可能是保存时混入了优化器状态
  • 若保存时用了 torch.compile(),加载后需重新 compile,state_dict 不包含编译后的图信息

保存/加载优化器状态也用 state_dict,但必须和模型分开管理

优化器也有自己的 state_dict(含动量、累积梯度等),常和模型权重一起保存用于断点续训。但它和模型的 state_dict 是两个独立字典,不能合并或混用。

Hypotenuse AI
Hypotenuse AI

AI写作助手和文本生成器,根据关键词生成原创的、有洞察力的文章

下载

典型做法是打包成一个 Python 字典再保存:

torch.save({
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
    'loss': loss,
}, path)

加载时也需分别取出、分别加载:

checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
  • optimizer.state_dict() 包含张量,保存路径必须支持二进制写入(不能用文本模式打开)
  • 不同优化器的 state_dict 结构差异大,比如 AdamWexp_avg_sq,而 SGD 没有;切换优化器时不能复用旧的 optimizer_state_dict
  • 如果用了 torch.cuda.amp.GradScaler,它的 state_dict 也要一并保存,否则混合精度训练会中断

跨设备加载要注意 map_location,否则 CPU 加载 GPU 模型会卡死

在 GPU 上训练保存的模型,如果直接在 CPU 环境下加载,PyTorch 默认尝试把所有张量映射回原设备(即 CUDA),结果报 RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False

解决方法是在 torch.load() 里加 map_location 参数,显式指定目标设备:

  • CPU 加载:用 torch.load(path, map_location='cpu')
  • 指定某张 GPU:用 torch.load(path, map_location='cuda:1')
  • 自动适配当前设备:用 torch.load(path, map_location=torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

这个参数只影响 torch.load() 阶段,不影响后续 model.to(device) —— 但如果你先 load_state_dict()to(device),中间参数会经历一次设备拷贝,多一次内存开销。

模型权重序列化这件事,表面是读写文件,实际核心约束在「结构一致性」和「设备上下文」。哪怕路径、后缀、函数名都对,只要模型类定义稍有改动,或加载时设备没对齐,就会静默失效——不是报错,而是预测结果完全不对。

热门AI工具

更多
DeepSeek
DeepSeek

幻方量化公司旗下的开源大模型平台

豆包大模型
豆包大模型

字节跳动自主研发的一系列大型语言模型

WorkBuddy
WorkBuddy

腾讯云推出的AI原生桌面智能体工作台

腾讯元宝
腾讯元宝

腾讯混元平台推出的AI助手

文心一言
文心一言

文心一言是百度开发的AI聊天机器人,通过对话可以生成各种形式的内容。

讯飞写作
讯飞写作

基于讯飞星火大模型的AI写作工具,可以快速生成新闻稿件、品宣文案、工作总结、心得体会等各种文文稿

即梦AI
即梦AI

一站式AI创作平台,免费AI图片和视频生成。

ChatGPT
ChatGPT

最最强大的AI聊天机器人程序,ChatGPT不单是聊天机器人,还能进行撰写邮件、视频脚本、文案、翻译、代码等任务。

相关专题

更多
python开发工具
python开发工具

php中文网为大家提供各种python开发工具,好的开发工具,可帮助开发者攻克编程学习中的基础障碍,理解每一行源代码在程序执行时在计算机中的过程。php中文网还为大家带来python相关课程以及相关文章等内容,供大家免费下载使用。

781

2023.06.15

python打包成可执行文件
python打包成可执行文件

本专题为大家带来python打包成可执行文件相关的文章,大家可以免费的下载体验。

697

2023.07.20

python能做什么
python能做什么

python能做的有:可用于开发基于控制台的应用程序、多媒体部分开发、用于开发基于Web的应用程序、使用python处理数据、系统编程等等。本专题为大家提供python相关的各种文章、以及下载和课程。

784

2023.07.25

format在python中的用法
format在python中的用法

Python中的format是一种字符串格式化方法,用于将变量或值插入到字符串中的占位符位置。通过format方法,我们可以动态地构建字符串,使其包含不同值。php中文网给大家带来了相关的教程以及文章,欢迎大家前来阅读学习。

890

2023.07.31

python教程
python教程

Python已成为一门网红语言,即使是在非编程开发者当中,也掀起了一股学习的热潮。本专题为大家带来python教程的相关文章,大家可以免费体验学习。

1564

2023.08.03

python环境变量的配置
python环境变量的配置

Python是一种流行的编程语言,被广泛用于软件开发、数据分析和科学计算等领域。在安装Python之后,我们需要配置环境变量,以便在任何位置都能够访问Python的可执行文件。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

598

2023.08.04

python eval
python eval

eval函数是Python中一个非常强大的函数,它可以将字符串作为Python代码进行执行,实现动态编程的效果。然而,由于其潜在的安全风险和性能问题,需要谨慎使用。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

591

2023.08.04

scratch和python区别
scratch和python区别

scratch和python的区别:1、scratch是一种专为初学者设计的图形化编程语言,python是一种文本编程语言;2、scratch使用的是基于积木的编程语法,python采用更加传统的文本编程语法等等。本专题为大家提供scratch和python相关的文章、下载、课程内容,供大家免费下载体验。

759

2023.08.11

bootstrap安装教程
bootstrap安装教程

本专题整合了bootstrap安装相关教程,阅读专题下面的文章了解更多详细操作教程。

22

2026.03.18

热门下载

更多
网站特效
/
网站源码
/
网站素材
/
前端模板

精品课程

更多
相关推荐
/
热门推荐
/
最新课程
最新Python教程 从入门到精通
最新Python教程 从入门到精通

共4课时 | 22.5万人学习

Django 教程
Django 教程

共28课时 | 5.1万人学习

SciPy 教程
SciPy 教程

共10课时 | 2万人学习

关于我们 免责申明 举报中心 意见反馈 讲师合作 广告合作 最新更新
php中文网:公益在线php培训,帮助PHP学习者快速成长!
关注服务号 技术交流群
PHP中文网订阅号
每天精选资源文章推送

Copyright 2014-2026 https://www.php.cn/ All Rights Reserved | php.cn | 湘ICP备2023035733号