pytorch学习笔记-----网络模型的保存和读取

模型的保存是调用torch的save方法，一共有两种保存方式

import torch
import torchvision
#模型的保存
#使用现有的网络模型-》假装得到一个自己的模型
vgg16 = torchvision.models.vgg16()

#保存方式1，这种方式不仅保存了网络训练的参数，还保存了网络的结构
torch.save(vgg16,"vgg16_method1.pth")

#保存方式2，这种方式仅仅保存了网络训练的参数,是通过字典的方式进行保存
torch.save(vgg16.state_dict(),'vgg16_method2.pth')

模型的加载是调用torch下面的load方法， 和模型的保存一样模型的加载也是有俩种方法。第一种比较简单：直接加载模型就可以。第二种需要先把网络创建成对象，然后在加载报讯的权重参数。

 

import torch
from model_save import *#使用方式1，需要在加载模型的时候引入保存的网络


#模型的加载
#这种方式加载不仅加载了模型还加载了训练好的参数---->方式1
import torchvision

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

#加载模型，初始化参数的那种，--------》方式2
vgg16 =torchvision.models.vgg16()

#将参数装到模型之中
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)

#这种方式只加载了模型的参数
# model = torch.load("vgg16_method2.pth")
# print(model)