【学习笔记】【Pytorch】十五、网络模型的保存与读取

Mr庞.

已于 2023-01-18 11:12:05 修改

阅读量534

点赞数

CC 4.0 BY-SA版权

分类专栏： Pytorch 文章标签： pytorch 学习网络

于 2023-01-16 22:16:50 首次发布

本文链接：https://blog.csdn.net/qq_45772756/article/details/128709901

文章介绍了Pytorch中两种保存和读取网络模型的方法。保存模型可以是整个模型结构和参数，或者仅保存state_dict。在读取时，需要注意加载模型时需提供相应的模型类定义。还特别提到了加载自定义模型时避免缺少类定义的错误处理方式。

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

【学习笔记】【Pytorch】十五、网络模型的保存与读取

一、网络模型的保存
二、网络模型的读取

一、网络模型的保存

Pytorch提供了两种方式进行保存模型。

import torch
import torchvision
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d

vgg16 = torchvision.models.vgg16(pretrained=False)

# 保存方式1：模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")  # 保存模型结构及参数

# 保存方式2：模型参数，保存成字典的形式（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")


# 陷阱1：方式1保存模型，陷阱在加载处
class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()  # 初始化父类属性
        self.model1 = Sequential(
            Conv2d(3, 32, 5, stride=1, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5