pytorch模型保存

最新推荐文章于 2025-06-28 10:37:12 发布

tang-0203

最新推荐文章于 2025-06-28 10:37:12 发布

阅读量2.7k

点赞数 2

CC 4.0 BY-SA版权

分类专栏： Pytorch学习文章标签： pytorch 模型保存模型保存

本文链接：https://blog.youkuaiyun.com/tsq292978891/article/details/79434892

Pytorch学习专栏收录该内容

24 篇文章

订阅专栏

本文介绍PyTorch中两种常见的模型保存与加载方法：一种仅保存模型参数，适用于模型迁移；另一种同时保存模型结构与参数，方便直接加载使用。文章通过具体示例展示了如何操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

#pytorch模型保存方式：
pytorch官网手册

torchvision.models中的下载下来的pretrained模型，比如：alexnet-owt-4df8aa71.pth，是将模型的参数，按照有序字典的方式保存。类型为：collections.OrderedDict

###方式一：仅保存和加载模型参数（推荐的方式）

1 . 保存模型参数：

import torch
torch.save(model.state_dict(), 'save_path_name.pth') # 后缀不重要(torchvision.models下载下来的模型参数后缀为.pth)

2 . 加载模型参数：

import torch
import torch.nn as nn
model.load_state_dict(torch.load('save_path_name.pth'), strict=True) # strict=True表示键值要严格匹配

例子：

import torch
import torch.nn as nn
import torchvision
import AlexNet_Train_Val

model = AlexNet_Train_Val.ModifiedAlexNet(2)
model.load_state_dict(torch.load('./models/AlexNet1.pth'), strict=True) # AlexNet1.pth为之前保存的模型参数文件

###方式二：保存和加载整个模型(模型结构和模型参数)
1 . 保存模型：

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

2 . 加载模型：

import torch
import torch.nn as nn
model = torch.load('save_path_name.pth')

额外的参考：https://zhuanlan.zhihu.com/p/82038049