PyTorch保存和加载模型的两种方式

与Tensorflow、Keras等框架一样，Pytorch也提供了两种保存模型的方式：

1. 只保存模型参数

2. 保存完整模型

下面我们依次对这两种方式进行实现，以以下多层感知机模型为例：

import torch.nn as nn
import torch

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.pth')

加载模型参数：

# 定义模型结构
net_clone = create_net()
# 加载模型参数
net_clone.load_state_dict(torch.load('net_parameter.pth'))

而后我们就可以微调模型或调用模型进行预测。

2、保存完整模型

保存完整模型指既保存模型结构又保存模型参数：

torch.save(net, 'net_model.pth')

加载模型：

net_loaded = torch.load('net_model.pth')

从上图我们可以看到，只保存模型参数的文件大小为4kb，而保存整个模型的文件大小为12kb，对于一些复杂的模型而言，保存整个模型可能需要占用大量的存储空间。所以推荐保存模型参数，在实际预测时再重新定义模型结构即可。