在深度学习项目中，模型的保存与加载是一个不可忽视的环节。它不仅涉及到模型权重的持久化，还包括了训练状态的保存和恢复。此外，在训练大型模型或者进行长时间训练时，我们经常需要使用到断点续训的技巧，以应对计算资源的限制和不可预见的中断。本文将深入探讨PyTorch中模型保存与加载的方法，并分享一些实用的断点续训技巧。
模型保存与加载
在PyTorch中，模型的保存与加载主要依赖于torch.save和torch.load两个函数。这些函数可以保存模型的状态字典（state_dict）或者整个模型对象，并能够在之后的训练中加载和恢复。
保存模型
保存状态字典
模型的状态字典包含了模型中每一层的参数，是一种轻量级的保存方式。通常推荐使用这种方法，因为它不保存模型的计算图结构，节省空间且更加灵活。
torch.save(model.state_dict(), 'model_weights.pth')
保存完整模型
如果需要保存模型的完整结构和权重，可以直接保存模型对象。
torch.save(model, 'model_complete.pth')
加载模型
加载状态字典
加载状态字典时，需要先创建一个与保存时相同结构的模型实例，然后使用load_state_dict方法。
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load('model_weights.pth'))
加载完整模型
直接加载模型对象是一种更加简便的方式，但需要注意，这种方式会同时加载模型的结构和权重。
model = torch.load('model_complete.pth')
断点续训技巧
在实际训练过程中，可能会遇到各种中断情况，如计算资源的分配、意外断电等。为了应对这些情况，我们需要掌握一些断点续训的技巧。
保存训练状态
除了保存模型权重，我们还应该保存足够的训练状态信息，以便于后续从断点处继续训练。这些信息通常包括：
- 当前的训练轮次（epoch）
- 训练数据的批次索引（batch index）
- 优化器的状态（optimizer state）
- 学习率调度器的状态（scheduler state）
# 假设我们有一个包含模型、优化器和学习率调度器的字典 `state`
state = {
'epoch': current_epoch,
'batch_index': current_batch_index,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict() if scheduler else None,
'model_state_dict': model.state_dict()
}
# 保存训练状态
torch.save(state, 'training_state.pth')
从断点续训
当需要从断点继续训练时，我们首先加载保存的训练状态，然后恢复模型、优化器和学习率调度器的状态。
# 加载训练状态
state = torch.load('training_state.pth')
# 创建模型和优化器实例
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*optimizer_args, **optimizer_kwargs)
# 恢复训练状态
model.load_state_dict(state['model_state_dict'])
optimizer.load_state_dict(state['optimizer'])
if state['scheduler'] is not None:
scheduler = TheSchedulerClass(*scheduler_args, **scheduler_kwargs)
scheduler.load_state_dict(state['scheduler'])
# 更新当前的训练轮次和批次索引
current_epoch = state['epoch']
current_batch_index = state['batch_index']
# 继续训练
print(f'Resuming training from epoch {current_epoch}, batch {current_batch_index + 1}...')
结论
模型的保存与加载以及断点续训是深度学习工程实践中的重要技能。通过本文的介绍，我们学习了PyTorch中相关的操作方法和技巧。合理地保存模型和训练状态，可以帮助我们在面对训练中断时，快速恢复训练过程，节省时间和计算资源。掌握这些技巧，将有助于我们更加高效地进行深度学习项目的开发和迭代。
