pytorch模型保存及加载
代码
import torch
import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False)# 1. save model 1 保存模型结构及模型参数
torch.save(vgg16, './vgg16_save1.model')# 2. save model 2 只保存模型参数 比第一种保存方法保存的文件要小
torch.save(vgg16.state_dict(), './vgg16_save2.model')# 3. load model 1
vgg16_load1 = torch.load('./vgg16_save1.model')
print(vgg16_load1) # 打印的是模型网络结构# 3. load model 2
vgg16_load2 = torch.load('./vgg16_save2.model')
print(vgg16_load2) # 打印的是模型参数
# 将参数导入到网络
vgg16.load_state_dict(vgg16_load2)
print(vgg16)# 5. 保存模型方式1的陷阱
# 当用方法1导入模型的时候,模型结构是要已知的
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
# class MySeq2(nn.Module):
# def __init__(self):
# super(MySeq2, self).__init__()
# self.model1 = Sequential(Conv2d(3, 32, kernel_size=5, stride=1, padding=2),
# MaxPool2d(2),
# Conv2d(32, 32, kernel_size=5, stride=1, padding=2),
# MaxPool2d(2),
# Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
# MaxPool2d(2),
# Flatten(),
# Linear(1024, 64),
# Linear(64, 10)
# )
#
# def forward(self, x):
# x = self.model1(x)
# return x
# myseq2 = MySeq2()
# torch.save(myseq2, 'myseq_self.model')
# 当用方法1导入模型的时候,模型结构是要已知的 否则就会报下面的错误 可以在代码里重新定义 但一般都是写在另一个单独的文件里面 比如上面注释的模型结构是前面已经写在p19_nn_seq 文件里面的,执行了模型保存
# AttributeError: Can't get attribute 'MySeq2' on <module '__main__' from 'C:/工作文档/learn_pytorch/p23_save_load_model.py'>
from p19_nn_seq import *
myseq2 = torch.load('myseq_self.model')
print(myseq2)
执行结果
只打印模型参数
打印模型结构,在调试模式下 可以在feature–保护属性–models–0–weight下看到模型参数
自己写过的模型文件保存后加载