网络模型的保存与读取主要是有两种方式
一:保存网络模型的结构
二:获取网络模型的状态并将其网络模型中的参数保存为字典
方式一:
保存模型:
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")