算力平台:
检查点
在使用 Accelerate 训练 PyTorch 模型时,你可能经常需要保存和继续训练状态。这样做需要保存和加载模型、优化器、RNG 生成器和 GradScaler。Accelerate 内部提供了两个便捷函数来快速实现这一点:
- 使用 [
~Accelerator.save_state
] 将上述所有内容保存到一个文件夹位置 - 使用 [
~Accelerator.load_state
] 从之前的save_state
加载所有存储的内容
为了进一步自定义通过 [~Accelerator.save_state
] 保存状态的位置和方式,可以使用 [~utils.ProjectConfiguration
] 类。例如,如果启用了 automatic_checkpoint_naming
,每个保存的检查点将位于 Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}
。
需要注意的是,这些状态应来自同一个训练脚本,不应来自两个不同的脚本。
- 通过使用 [
~Accelerator.register_for_checkpointing
],你可以注册自定义对象以自动存储或从上述两个函数加载,只要该对象具有state_dict
和load_state_dict
功能。这可以包括学习率调度器等对象。
以下是一个简短的示例,演示如何在训练过程中使用检查点保存和重新加载状态:
python
from accelerate import Accelerator
import torch
accelerator = Accelerator(project_dir="my/save/path")
my_scheduler = torch.optim.lr_scheduler.StepLR(my_optimizer, step_size=1, gamma=0.99)
my_model, my_optimizer, my_training_dataloader = accelerator.prepare(my_model, my_optimizer, my_training_dataloader)
# Register the LR scheduler
accelerator.register_for_checkpointing(my_scheduler)
# Save the starting state
accelerator.save_state()
device = accelerator.device
my_model.to(device)
# Perform training
for epoch in range(num_epochs):
for batch in my_training_dataloader:
my_optimizer.zero_grad()
inputs, targets = batch
inputs = inputs.to(device)
targets = targets.to(device)
outputs = my_model(inputs)
loss = my_loss_function(outputs, targets)
accelerator.backward(loss)
my_optimizer.step()
my_scheduler.step()
# Restore the previous state
accelerator.load_state("my/save/path/checkpointing/checkpoint_0")
恢复 DataLoader 的状态
在从检查点恢复后,如果在 epoch 的中途保存了状态,可能还需要从活动的 DataLoader
的某个特定点恢复。你可以使用 [~Accelerator.skip_first_batches
] 来实现这一点。
python
from accelerate import Accelerator
accelerator = Accelerator(project_dir="my/save/path")
train_dataloader = accelerator.prepare(train_dataloader)
accelerator.load_state("my_state")
# Assume the checkpoint was saved 100 steps into the epoch
skipped_dataloader = accelerator.skip_first_batches(train_dataloader, 100)
# After the first iteration, go back to `train_dataloader`
# First epoch
for batch in skipped_dataloader:
# Do something
pass
# Second epoch
for batch in train_dataloader:
# Do something
pass