pytorch学习笔记11-网络模型的读取与保存
目录
- 模型的保存
-
- 方式一
- 方式二
- 模型的读取
-
- 对应方式一
- 对应方式二
- 需要注意的问题
模型的保存
先引入vgg16模型（没有经过预训练的）
import torchvision
vgg16_false = torchvision.models.vgg16(pretained=False)
方式一
保存网络模型的结构和其中的参数
torch.save(vgg16_false, "vgg16_method1.pth")
在左侧可以看到已经保存的文件
方式二
把模型的参数保存成字典形式，不保存网络结构，官方推荐的保存方式，因为这种保存方式占用空间小
torch.save(vgg16_false.state_dict(), "vgg16_method2.pth")
在terminal中输入dir查看文件，可以看到方式二比方式一要小一些。
方式1用pth后缀,方式2用pkl后缀,便于区分
模型的读取
对应方式一
import torch
model = torch.load("vgg16_method1.pth")
print(model)
输出结果如下，可以debug对比一下方式一保存时的参数，都是一样的。
VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace=True)
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(inplace=True)
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(inplace=True)
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace=True)
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace=True)
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace=True)
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(inplace=True)
(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(inplace=True)
(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(inplace=True)
(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)
对应方式二
model2 = torch.load("vgg16_method2.pth")
print(model2)
输出是字典
如果想要输出网络结构需要这么写
vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_false.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16_false)
需要注意的问题
对自己创建的网络模型使用方式一保存时，读取时会出现问题。
比如自己创建一个模型
import torch
from torch import nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
def forward(self, x):
output = self.conv1(x)
return output
model = Model()
torch.save(model, "model_method1.pth")
按方式一的读取方法
import torch
model = torch.load("model_method1.pth")
print(model)
输出会报错
这时需要把自己创建的类导入到当前文件中，不需要进行实例化
import torch
from torch import nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
def forward(self, x):
output = self.conv1(x)
return output
# model = Model()
model = torch.load("model_method1.pth")
print(model)
或者
import torch
from torch import nn
from P26_model_save import *
# class Model(nn.Module):
# def __init__(self):
# super(Model, self).__init__()
# self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
#
# def forward(self, x):
# output = self.conv1(x)
# return output
# model = Model()
model = torch.load("model_method1.pth")
print(model)
这时程序可以正常输出
