跳到主要内容

PyTorch 序列化与保存

在深度学习中，模型的训练通常需要大量的时间和计算资源。为了避免每次需要时都重新训练模型，我们可以将训练好的模型保存下来，以便在以后的任务中直接加载和使用。PyTorch提供了多种方法来序列化和保存模型，本文将详细介绍这些方法。

什么是序列化？​

序列化是指将数据结构或对象状态转换为可以存储或传输的格式的过程。在PyTorch中，序列化通常指的是将模型的状态（包括权重和优化器状态）保存到文件中，以便在需要时可以重新加载。

保存和加载模型​

保存模型​

在PyTorch中，我们可以使用 torch.save() 函数来保存模型。通常，我们会保存模型的 state_dict，这是一个包含模型所有参数（权重和偏置）的字典。

python
import torch
import torch.nn as nn

# 定义一个简单的模型
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 实例化模型
model = SimpleModel()

# 保存模型的 state_dict
torch.save(model.state_dict(), 'model.pth')

加载模型​

要加载保存的模型，我们首先需要实例化一个与保存时相同的模型结构，然后使用 torch.load() 加载 state_dict，最后使用 load_state_dict() 方法将参数加载到模型中。

python
# 实例化模型
model = SimpleModel()

# 加载模型的 state_dict
model.load_state_dict(torch.load('model.pth'))

# 将模型设置为评估模式
model.eval()
备注

在加载模型后，记得调用 model.eval() 将模型设置为评估模式。这是因为某些层（如 Dropout 和 BatchNorm）在训练和评估时的行为是不同的。

保存和加载整个模型​

除了保存 state_dict，我们还可以保存整个模型。这种方法更加简单，但不够灵活，因为它依赖于保存时的模型类和代码。

python
# 保存整个模型
torch.save(model, 'model_entire.pth')

# 加载整个模型
model = torch.load('model_entire.pth')
model.eval()
警告

保存整个模型可能会导致代码的兼容性问题，尤其是在模型类定义发生变化时。因此，推荐使用 state_dict 的方法来保存和加载模型。

保存和加载检查点​

在训练过程中，我们通常希望保存模型的检查点（checkpoint），以便在训练中断时可以从中断处继续训练。检查点通常包括模型的 state_dict、优化器的 state_dict 以及当前的 epoch 和损失等信息。

python
# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

# 保存检查点
checkpoint = {
    'epoch': 10,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': 0.02,
}
torch.save(checkpoint, 'checkpoint.pth')

# 加载检查点
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

# 继续训练
model.train()

实际应用场景​

场景1：模型部署​

在模型部署时，我们通常需要将训练好的模型保存下来，然后在生产环境中加载并使用。使用 state_dict 的方法可以确保模型的灵活性和兼容性。

场景2：模型共享​

当你需要与团队成员或社区共享模型时，保存模型的 state_dict 是一个很好的选择。这样，其他人可以在他们的代码中加载模型，而不需要完全相同的模型类定义。

场景3：训练中断恢复​

在长时间的训练任务中，保存检查点可以帮助你在训练中断时从中断处恢复训练，而不需要从头开始。

总结​

在本文中，我们介绍了如何在PyTorch中序列化和保存模型。我们讨论了保存和加载模型的 state_dict、保存整个模型以及保存和加载检查点的方法。我们还探讨了这些方法在实际应用中的场景。

附加资源与练习​

  • 练习1：尝试在一个简单的模型上实现保存和加载 state_dict，并验证加载后的模型是否与保存前的模型一致。
  • 练习2：在一个训练任务中实现检查点的保存和加载，并在训练中断后恢复训练。
  • 资源：阅读 PyTorch官方文档 以了解更多关于序列化和保存的细节。

通过掌握这些技能，你将能够更高效地管理和使用你的深度学习模型。