网络模型保存与读取

最新推荐文章于 2024-06-13 23:55:30 发布

YLTommi

最新推荐文章于 2024-06-13 23:55:30 发布

阅读量74

点赞数

文章标签： python pytorch 开发语言

本文链接：https://blog.csdn.net/yangleileileilei/article/details/138974764

版权

# 1.网络模型保存（方式一）
# 方式一：模型结构+模型参数
import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)
torch.save(vgg16, "D:/PyCharm/vgg16model/vgg16_mothod1.pth")  # 方式一：模型结构+模型参数
print(vgg16)

# 2.网络模型导入 （方式一）
# 方式一：模型结构+模型参数 导入
model1 = torch.load("D:/PyCharm/vgg16model/vgg16_mothod1.pth")
print(model1)

# 3.网络模型保存 （方式二）
# 只保存模型参数，不保存模型结构 （官方推荐）
torch.save(vgg16.state_dict(), "D:/PyCharm/vgg16model/vgg16_mothod2.pth")  # 只保存模型参数，不保存模型结构 （官方推荐）
print(vgg16)

# 4.网络模型导入 （方式二）
model2 = torch.load("D:/PyCharm/vgg16model/vgg16_mothod2.pth")  # 只导入模型参数
print(model2)

vgg16.load_state_dict(torch.load("D:/PyCharm/vgg16model/vgg16_mothod2.pth"))  # 将模型参数导入到模型结构中
print(vgg16)


# 5.创建模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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


net = Net()
torch.save(net, "D:/PyCharm/vgg16model/net_mothod1.pth")

# 6.加载模型
net1 = torch.load("D:/PyCharm/vgg16model/net_mothod1.pth")
print(net1)