Accelerate documentation
チェックポイント
You are viewing main version, which requires installation from source. If you'd like
regular pip install, checkout the latest stable version (v1.15.0).
チェックポイント
AccelerateでPyTorchモデルをトレーニングする際、トレーニングの状態を保存して再開したい場面はよくあります。これには、モデル、オプティマイザー、RNGジェネレーター、GradScalerの保存と読み込みが必要です。Accelerateには、これを素早く実現するための2つの便利な関数が用意されています。
- 上記で説明したすべての内容をフォルダに保存するには、save_state()を使用します
- 以前に
save_stateで保存されたすべてのデータを読み込むには、load_state()を使用します。
save_state() による状態の保存先と保存方法のさらなるカスタマイズには、ProjectConfiguration クラスを使用できます。例えば、automatic_checkpoint_naming が有効な場合、各保存されたチェックポイントは Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number} に配置されます。
これらの状態が同じトレーニングスクリプトから生成されるものであることが前提であり、2つの別々のスクリプトから生成されたものであるべきではない点に注意が必要です。
state_dictおよびload_state_dictの機能を持つオブジェクトであれば、register_for_checkpointing() を使用してカスタムオブジェクトを登録し、前述の2つの関数によって自動的に保存または読み込むことができます。これには、学習率スケジューラーなどのオブジェクトが含まれます。
以下は、トレーニング中に状態を保存して読み戻すためにチェックポイントを使用する簡単な例です：
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の状態の復元
チェックポイントからの再開後、エポックの途中で状態が保存されていた場合、アクティブなDataLoaderの特定のポイントから再開することも望ましい場合があります。skip_first_batches()を使用することで、これを実現できます。
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