一、引言
我们今天来看一下模型的保存与加载~
我们平时在神经网络的训练时间可能会很长，为了在每次使用模型时避免高代价的重复训练，我们就需要将模型序列化到磁盘中，使用的时候反序列化到内存中。
PyTorch提供了两种主要的方法来保存和加载模型，分别是直接序列化模型对象和存储模型的网络参数。
二、直接序列化模型对象
直接序列化模型对象：方法使用torch.save()函数将整个模型对象保存为一个文件，然后使用torch.load()函数将其加载回内存。这种方法可以方便地保存和加载整个模型，包括其结构、参数以及优化器等信息。
import torch
import torch.nn as nn
import pickleclass 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 outputdef 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 的问题，特别是在加载模型时需要注意 :
保存和加载设备一致性:
- 当你在 GPU 上训练了一个模型，并使用
torch.save()保存了该模型的状态字典（state_dict），然后尝试在一个没有 GPU 的环境中加载该模型时，会引发错误，因为 PyTorch 期望在相同的设备上执行操作。- 为了解决这个问题，你可以在没有 GPU 的机器上保存整个模型（而不是仅保存
state_dict），这样 PyTorch 会将权重数据移动到 CPU 上，并且在加载时不会引发错误。移动模型到 CPU:
- 如果你在 GPU 上保存了模型的
state_dict，并且想在 CPU 上加载它，你需要确保在加载state_dict之前将模型移动到 CPU。这可以通过调用模型的to('cpu')方法来实现。移动模型到 GPU:
- 如果你在 CPU 上保存了模型的
state_dict，并且想在 GPU 上加载它，你需要确保在加载state_dict之前将模型移动到 GPU。这可以通过调用模型的to(device)方法来实现，其中device是一个包含 CUDA 信息的对象（如果 GPU 可用）。
三、存储模型的网络参数
import torch
import torch.nn as nn
import torch.optim as optimclass 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 outputdef 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，这是控制参数更新步长的重要超参数。
