一、示例
这里是一个示例,展示如何保存和加载模型、优化器的状态以恢复训练:
保存与加载模型和优化器状态的示例
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
),这些数据同样需要额外的存储空间。
- 对于像
-
学习率调度器的状态:
- 如果使用了学习率调度器,优化器状态文件中也可能包含与调度器相关的状态信息(例如,学习率的历史变化)。
虽然优化器文件不保存权重参数本身,但因为每个参数可能都有多个与其相关的状态(如动量、一阶和二阶矩等),这些状态值累积起来可能会使优化器的文件比权重文件还大。
- 与参数个数成比例:
优化器的状态大小通常与模型参数的数量成比例。例如,如果模型有 1 亿个参数,而 Adam 优化器对每个参数存储两个额外的状态变量,那么优化器的状态文件可能会是模型权重文件大小的 2 倍(2 个状态变量)。
举例说明
假设你有一个 100 MB 大小的模型权重文件,如果你使用 Adam 优化器,那么优化器状态文件可能会有 200 MB 大小(200 MB 的状态信息,如动量和平方梯度)。
简单总结:
优化器的状态文件通常比权重文件大,主要原因是它保存了用于优化过程的额外信息,如动量、累积梯度等,而不包括模型的权重参数。