pytorch模型保存及加载参数恢复训练的例子

静心问道

已于 2024-09-02 16:56:20 修改

阅读量447

点赞数 7

分类专栏：深度学习文章标签： pytorch python 人工智能

于 2024-09-02 16:28:50 首次发布

本文链接：https://blog.csdn.net/qq_39698985/article/details/141823143

版权

深度学习专栏收录该内容

15 篇文章 0 订阅

订阅专栏

一、示例

这里是一个示例，展示如何保存和加载模型、优化器的状态以恢复训练：

保存与加载模型和优化器状态的示例

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

# 定义一个简单的模型
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()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 定义一个损失函数
criterion = nn.MSELoss()

# 模拟一些训练数据
inputs = torch.randn(10, 10)
targets = torch.randn(10, 1)

# 模拟训练步骤
for epoch in range(5):  # 训练5个epoch
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    print(f"Epoch [{epoch+1}/5], Loss: {loss.item()}")

# 保存模型和优化器的状态
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss.item()
}
torch.save(checkpoint, 'checkpoint.pth')

print("\nTraining complete, model and optimizer state saved.\n")

# 恢复训练的过程
# 加载保存的状态
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1  # 从保存时的下一个epoch开始
loss = checkpoint['loss']

print(f"Resuming training from epoch {start_epoch} with loss {loss:.4f}\n")

# 模拟恢复训练后继续训练
for epoch in range(start_epoch, 10):  # 继续训练到第10个epoch
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    print(f"Epoch [{epoch+1}/10], Loss: {loss.item()}")

示例解释

模型定义与训练
- 定义了一个简单的全连接模型 SimpleModel。
- 使用随机数据模拟了5个epoch的训练过程。
- 每个epoch中，模型计算输出，损失函数计算误差，优化器调整模型参数。
保存模型与优化器状态
- 使用 torch.save() 保存模型的状态字典、优化器的状态字典、当前的 epoch 和 loss。
- 保存后的状态被存储在一个 checkpoint.pth 文件中。
加载模型与恢复训练
- 使用 torch.load() 加载模型和优化器的状态。
- 加载后从保存时的 epoch 继续训练，确保恢复训练不重复已经完成的步骤。
继续训练
- 从保存的 epoch 开始，继续训练到第10个epoch。

输出

程序会在前5个epoch中训练模型，保存模型和优化器的状态。然后，加载保存的状态，并从第6个epoch开始继续训练，直到第10个epoch。

二、疑问1

不保存model.parameters()，那么优化器恢复训练时为啥不需要再次传递model.parameters()？

在恢复训练时，优化器不需要再次传递 model.parameters() 是因为优化器已经持有对模型参数的引用。在初始化优化器时，例如通过 optimizer = optim.SGD(model.parameters(), lr=0.01)，优化器会保留对模型参数的引用，以便它能够在训练过程中更新这些参数。

工作原理

优化器初始化时绑定模型参数：
- 当你调用 optimizer = optim.SGD(model.parameters(), lr=0.01) 时，优化器会获取模型参数的迭代器，并将这些参数存储在它的内部数据结构中。优化器不复制这些参数，而是保留它们的引用。
优化器状态的保存：
- 当保存优化器的状态（如 optimizer.state_dict()）时，保存的是优化器内部的动量、累积梯度等信息。这些信息依赖于模型参数，但并不直接包含模型权重本身。
恢复训练时的优化器状态加载：
- 在恢复训练时，首先加载模型的权重（model.load_state_dict()），然后加载优化器的状态（optimizer.load_state_dict()）。
- 因为优化器仍然持有对 model.parameters() 的引用，所以它可以继续更新这些参数。优化器只需要知道如何继续进行参数更新（通过保存的动量、累积梯度等信息），而不需要重新绑定模型参数。

为什么不需要再次传递 `model.parameters()`？

引用保持不变：优化器在初始化时已经持有模型参数的引用，而模型的权重加载（model.load_state_dict()）是直接对这些参数进行修改的。因此，优化器仍然指向同一组参数，不需要再次传递。
优化器只管理状态信息：优化器的状态信息（如动量、累积梯度等）与特定的模型参数相关。当你加载模型的权重后，优化器会继续基于当前模型权重来执行更新操作，而不需要重新获取模型参数。

总结

在恢复训练时，优化器不需要重新传递 model.parameters() 是因为它在初始化时已经绑定了模型参数，并且这些引用在模型权重加载后仍然有效。优化器只需要恢复它的状态，就可以继续操作绑定的模型参数。

你是对的，优化器的状态文件不包括模型的权重参数。因此，优化器文件比模型参数文件大的原因与保存的中间状态信息有关，而不是因为它包含了模型的参数。

二、疑问2

为什么优化器的文件比参数文件还大？

解释更新：

它保存了用于更新这些权重的各种状态信息。例如：

动量（Momentum）：
- 在优化算法中，像 SGD with momentum 或 Adam 这样的优化器需要存储每个参数的动量（或者称为梯度的累积值）。这些值独立于模型的权重参数存储。
二阶矩估计：
- 对于像 Adam 这样的自适应优化算法，它还需要存储每个参数的梯度的一阶和二阶移动平均值（exp_avg 和 exp_avg_sq），这些数据同样需要额外的存储空间。
学习率调度器的状态：
- 如果使用了学习率调度器，优化器状态文件中也可能包含与调度器相关的状态信息（例如，学习率的历史变化）。