PyTorch模型的保存加载

一、引言

我们今天来看一下模型的保存与加载~

我们平时在神经网络的训练时间可能会很长，为了在每次使用模型时避免高代价的重复训练，我们就需要将模型序列化到磁盘中，使用的时候反序列化到内存中。

PyTorch提供了两种主要的方法来保存和加载模型，分别是直接序列化模型对象和存储模型的网络参数。

二、直接序列化模型对象

直接序列化模型对象：方法使用torch.save()函数将整个模型对象保存为一个文件，然后使用torch.load()函数将其加载回内存。这种方法可以方便地保存和加载整个模型，包括其结构、参数以及优化器等信息。

import torch
import torch.nn as nn
import pickle

class Model(nn.Module):

    def __init__(self, input_size, output_size):

        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2, output_size)

    def forward(self, inputs):

        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output


def test01():

    model = Model(12, 2)

    # 第一个参数: 这是要保存的模型
    # 第二个参数: 这是模型保存的路径
    # 第三个参数: 指定了用于序列化和反序列化的模块
    # 第四个参数: 这是使用的pickle协议的版本,协议引入了二进制格式，提高了序列化数据的效率
    torch.save(model, 'model/test_model_save.pth', pickle_module=pickle, pickle_protocol=2)


def test02():

    # 第一个参数: 加载的路径
    # 第二个参数: 模型加载的设备
    # 第三个参数: 加载的模块
    model = torch.load('model/test_model_save.pth', map_location='cpu', pickle_module=pickle)

在使用 torch.save() 保存模型时，需要注意一些关于 CPU 和 GPU 的问题，特别是在加载模型时需要注意 :

1. 保存和加载设备一致性:

   • 当你在 GPU 上训练了一个模型，并使用 torch.save() 保存了该模型的状态字典（state_dict），然后尝试在一个没有 GPU 的环境中加载该模型时，会引发错误，因为 PyTorch 期望在相同的设备上执行操作。
   • 为了解决这个问题，你可以在没有 GPU 的机器上保存整个模型（而不是仅保存 state_dict），这样 PyTorch 会将权重数据移动到 CPU 上，并且在加载时不会引发错误。

2. 移动模型到 CPU:

   • 如果你在 GPU 上保存了模型的 state_dict，并且想在 CPU 上加载它，你需要确保在加载 state_dict 之前将模型移动到 CPU。这可以通过调用模型的 to('cpu') 方法来实现。

3. 移动模型到 GPU:

   • 如果你在 CPU 上保存了模型的 state_dict，并且想在 GPU 上加载它，你需要确保在加载 state_dict 之前将模型移动到 GPU。这可以通过调用模型的 to(device) 方法来实现，其中 device 是一个包含 CUDA 信息的对象（如果 GPU 可用）。

三、存储模型的网络参数

import torch
import torch.nn as nn
import torch.optim as optim


class Model(nn.Module):

    def __init__(self, input_size, output_size):

        super(Model, self).__init__()
        self.linear1 = nn.Linear(input_size, input_size * 2)
        self.linear2 = nn.Linear(input_size * 2, output_size)

    def forward(self, inputs):

        inputs = self.linear1(inputs)
        output = self.linear2(inputs)
        return output



def test01():

    model = Model(12, 2)
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    # 定义存储参数
    save_params = {
        'init_params': {
            'input_size': 12,
            'output_size': 2
        },
        'acc_score': 0.96,
        'avg_loss': 0.85,
        'iter_numbers': 100,
        'optim_params': optimizer.state_dict(),
        'model_params': model.state_dict()
    }

    # 存储模型参数
    torch.save(save_params, 'model/model_params.pth')


def test02():

    # 加载模型参数
    model_params = torch.load('model/model_params.pth')
    # 初始化模型
    model = Model(model_params['init_params']['input_size'], model_params['init_params']['output_size'])
    # 初始化优化器
    optimizer = optim.Adam(model.parameters())
    optimizer.load_state_dict(model_params['optim_params'])
    # 显示其他参数
    print('迭代次数:', model_params['iter_numbers'])
    print('准确率:', model_params['acc_score'])
    print('平均损失:', model_params['avg_loss'])

训练完成后，通常需要保存模型的参数值，以便用于后续的测试过程。使用torch.save()函数来保存模型的状态字典（state_dict），这个状态字典包含了模型的可学习参数（权重和偏置值）

optimizer = optim.Adam(model.parameters(), lr=0.01)

• 创建一个Adam优化器对象，在PyTorch中，优化器用于更新模型的参数以最小化损失函数。Adam是一种常用的优化算法，它结合了Momentum和RMSProp的优点，具有自适应学习率调整的特性。
• model.parameters()表示要优化的模型参数，即模型中所有可学习的权重和偏置值。lr=0.01表示学习率（learning rate）为0.01，这是控制参数更新步长的重要超参数。