网络模型的保存与读取

文章目录

一、模型的保存

方式1:torch.save(vgg16, "vgg16_method1.pth")

bash 复制代码
import torch
import torchvision.models

vgg16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg16, "vgg16_method1.pth")

如果运行报错:UserWarning: Arguments other than a weight enum or None for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing weights=None或者The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future

原因是在 PyTorch 的 torchvision 库中,从版本 0.13 开始,pretrained 参数已经被弃用,取而代之的是 weights 参数。这个改变是为了提供更丰富的预训练模型选择。当你尝试使用 vgg16(pretrained=False) 时,你收到了一个警告,告诉你 pretrained 参数已经不再被使用,并且建议你使用 weights 参数。

要解决这个问题,你应该使用 weights 参数来代替 pretrained。

修正代码:

bash 复制代码
import torch
import torchvision.models

vgg16 = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.DEFAULT)
torch.save(vgg16, "vgg16_method1.pth")

运行代码:

可以看到多了一个新文件vgg16_method1.pth

该方式1保存的网路模型不仅保存了网络模型的一种结构,它也保存了模型当中的一些参数

方式2:把模型的参数保存成字典(dict)形式

bash 复制代码
import torch
import torchvision.models

vgg16 = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.DEFAULT)
torch.save(vgg16.state_dict(),"vgg16_method2.pth")

运行结果:

方式1与方式2对比:

方式1保存的是模型的结构+模型的参数,方式2保存的只是模型的参数(官方推荐的保存方式)

官方推荐的原因是当保存一个大的模型时候,方式2所用的空间更小

我们可以查看一下两种保存方式的文件大小:

因为vgg这个模型本身就不大,所以文件大小差距并不明显,但方式2足足小了7kb!要是在大模型下节省空间这点会尤其明显。

二、文件的加载

代码:

bash 复制代码
import torch

model = torch.load("vgg16_method1.pth")
print(model)

运行结果:

通过将save与load的文件debug运行:

能够发现两者都是一样的,说明被完整加载出来。

通过上述步骤可以看到模型中的参数也一同保存下来了。

加载方式2保存的模型:

bash 复制代码
import torch

model = torch.load("vgg16_method2.pth")
print(model)

运行结果:

可以看到方式2形式是一个个字典形式.

方式2从字典形式想要恢复网络模型结构则需要:

bash 复制代码
import torch
import torchvision

# 创建一个VGG16模型实例,参数pretrained=False表示不加载预训练的权重。
vgg16 = torchvision.models.vgg16(pretrained=False)

# 加载之前保存的模型权重,这些权重保存在名为"vgg16_method2.pth"的文件中。
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))

# 打印出模型的结构,这样可以看到模型的各个层和参数。
print(vgg16)

运行结果:

可以看到把模型参数成功加载出来了

三、模型加载时容易犯的陷阱

保存一个自己写的网络模型:

bash 复制代码
import torch
import torchvision.models
from torch import nn


class Sen(nn.Module):
    def __init__(self):
        super(Sen, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x

sen = Sen()
torch.save(sen, "sen_method1.pth")

运行结果:

因为用方式1进行保存,故使用方式1的方法进行加载:

bash 复制代码
import torch
import torchvision

model = torch.load("sen_method1.pth")
print(model)

运行结果:

可以看到发生了报错,报错的意思是加载的时候没有找到Sen这个类

解决方法是将类复制到加载代码中:

bash 复制代码
import torch
import torchvision

class Sen(nn.Module):
    def __init__(self):
        super(Sen, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x

model = torch.load("sen_method1.pth")
print(model)

注意不需要写sen = Sen()这一代码

运行代码:

也就是用自己写的网络模型不同于现有的网络模型,需要进行导入才能正常加载出来!

或者也可以用import的方法加载自己写的网络模型,那么就不需要老是复制粘贴

通过from model_save import *加载:

bash 复制代码
import torch
from model_save import *

model = torch.load("sen_method1.pth")
print(model)

运行结果:

相关推荐
fyakm17 分钟前
RNN的注意力机制:原理与实现(代码示例)
rnn·深度学习·神经网络
金井PRATHAMA3 小时前
描述逻辑(Description Logic)对自然语言处理深层语义分析的影响与启示
人工智能·自然语言处理·知识图谱
Rock_yzh3 小时前
AI学习日记——参数的初始化
人工智能·python·深度学习·学习·机器学习
拆房老料4 小时前
Transformer推理优化全景:从模型架构到硬件底层的深度解析
深度学习·ai·自然语言处理·transformer
CiLerLinux5 小时前
第四十九章 ESP32S3 WiFi 路由实验
网络·人工智能·单片机·嵌入式硬件
青衫客365 小时前
基于 Python 构建的安全 gRPC 服务——TLS、mTLS 与 Casbin 授权实战
python·安全·微服务
-dzk-6 小时前
【3DGS复现】Autodl服务器复现3DGS《简单快速》《一次成功》《新手练习复现必备》
运维·服务器·python·计算机视觉·3d·三维重建·三维
七芒星20236 小时前
多目标识别YOLO :YOLOV3 原理
图像处理·人工智能·yolo·计算机视觉·目标跟踪·分类·聚类
Learn Beyond Limits7 小时前
Mean Normalization|均值归一化
人工智能·神经网络·算法·机器学习·均值算法·ai·吴恩达
byzy7 小时前
【论文笔记】VisionPAD: A Vision-Centric Pre-training Paradigm for Autonomous Driving
论文阅读·深度学习·计算机视觉·自动驾驶