与Tensorflow、Keras等框架一样,Pytorch也提供了两种保存模型的方式,这两种方式都是通过调用pickle序列化方法实现的:
1、只保存模型参数
2、保存完整模型
下面我们依次对这两种方式进行实现,以以下多层感知机模型为例:
def create_net():
net = nn.Sequential()
net.add_module('linear1', nn.Linear(15, 20))
net.add_module('relu1', nn.ReLU())
net.add_module('linear2', nn.Linear(20, 15))
net.add_module('relu2', nn.ReLU())
net.add_module('linear3', nn.Linear(15, 1))
net.add_module('sigmoid', nn.Sigmoid())
return net
net = create_net()
模型的基本信息如下:
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Linear-1 [-1, 20] 320
ReLU-2 [-1, 20] 0
Linear-3 [-1, 15] 315
ReLU-4 [-1, 15] 0
Linear-5 [-1, 1] 16
Sigmoid-6 [-1, 1] 0
================================================================
Total params: 651
Trainable params: 651
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.000057
Forward/backward pass size (MB): 0.000549
Params size (MB): 0.002483
Estimated Total Size (MB): 0.003090
----------------------------------------------------------------
None
1、只保存模型参数
保存模型参数:
torch.save(net.state_dict(), 'net_parameter.pkl')
加载模型参数:
# 定义模型结构
net_clone = create_net()
# 加载模型参数
net_clone.load_state_dict(torch.load('net_parameter.pkl'))
而后我们就可以微调模型或调用模型进行预测。
2、保存完整模型
保存完整模型指既保存模型结构又保存模型参数:
torch.save(net, 'net_model.pkl')
加载模型:
net_loaded = torch.load('net_model.pkl')
从上图我们可以看到,只保存模型参数的文件大小为4kb,而保存整个模型的文件大小为12kb,对于一些复杂的模型而言,保存整个模型可能需要占用大量的存储空间。所以推荐保存模型参数,在实际预测时再重新定义模型结构即可。