网络模型的保存与读取

网络模型的保存与读取主要是有两种方式
一：保存网络模型的结构
二：获取网络模型的状态并将其网络模型中的参数保存为字典

方式一：

保存模型：

import torch
import torchvision

#使用原始数据
vgg16 = torchvision.models.vgg16(pretrained=False)

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

加载模型：

import torch

#保存方式1，加载模型
model = torch.load("vgg16_methed1.pth")

方式二：

保存模型：

import torch
import torchvision

#使用原始数据
vgg16 = torchvision.models.vgg16(pretrained=False)

#保存方式2
#获取vgg16的状态并将其网络模型中的参数保存为字典
#模型参数（官方推荐）
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

加载模型：

import torch
import torchvision

#保存方式2，加载模型
#因为方式2保存的是字典类型的数据，所以要先创建vgg16模型，再加载参数
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))

可能出现的问题
如果该模型保存的时候有类，则在加载的时候需要引入该类

保存模型：

import torch
from torch import nn

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

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

test = Test()
torch.save(test, "test_method1.pth")

加载模型：

import torch
from model_load import *

model = torch.load("test_method1.pth")