Pytorch学习日记——网络模型的保存与读取

最新推荐文章于 2023-08-21 18:29:52 发布

快乐白给

最新推荐文章于 2023-08-21 18:29:52 发布

阅读量110

点赞数

分类专栏：学习文章标签： pytorch 学习深度学习 python Powered by 金山文档

本文链接：https://blog.csdn.net/Enjoylalala/article/details/128728042

版权

学习专栏收录该内容

13 篇文章 0 订阅

订阅专栏

文章介绍了如何在PyTorch中使用两种方法保存和加载网络模型，包括完整模型和仅保存state_dict的方式。还展示了加载预训练的vgg16模型以及自定义的Tudui网络模型的示例代码，并提到了加载自定义模型时的注意事项。

摘要由CSDN通过智能技术生成

学习视频——B站【小土堆】

网络模型的保存，代码

import torch
import torchvision
from torch import nn

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

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

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

# 陷阱
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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

tudui = Tudui()
torch.save(tudui, "tudui_method1.pth")

网络模型的读取，代码

import torch
import torchvision
from torch import nn

# 方式1，加载模型
# model = torch.load("vgg16_method1.pth")
# print(model)


# 方式二，加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# model2 = torch.load("vgg16_method2.pth")
# print(model2)
#print(vgg16)

# 陷阱1
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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


model3 = torch.load("tudui_method1.pth")
print(model3)
# 从模型当中import所有也可，如：
# from model_save import *
# model3 = torch.load("tudui_method1.pth")
# print(model3) # 其中model_save为模型所在文件