Acceleratorを使ったモデルの学習

PyTorch による効率的な AI モデルトレーニング

Dennis Lee

Data Engineer, Amazon

TrainerとAccelerator

AcceleratorとTrainerの使いやすさとカスタマイズ性を比較したグラフ。

カスタム学習ループ

Trainer：カスタム学習ループは不可
生成AIの高度なタスクでは2つのネットワークが必要になる場合がある

敵対的生成ネットワークの図。

¹ https://www.aitude.com/basics-of-generative-adversarial-network-model/

TrainerとAccelerator

AcceleratorとTrainerの使いやすさとカスタマイズ性を比較したグラフ。

基本的な学習ループの修正

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

    outputs = model(inputs)

    loss = outputs.loss

    loss.backward()

    optimizer.step()
    scheduler.step()

勾配をゼロに初期化
指定デバイスにデータを移動：.to(device)
順伝播の実行
クロスエントロピー損失の計算
逆伝播で勾配を計算
モデルパラメータと学習率を更新

Acceleratorオブジェクトの作成

Acceleratorは分散学習のインターフェースを提供します

from accelerate import Accelerator
accelerator = Accelerator(
    device_placement=True
)

device_placement（bool、デフォルトTrue）：デバイス配置を自動処理

モデルとオプティマイザの定義

事前学習済みモデルの読み込み

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-cased", return_dict=True)

Adamでモデルパラメータを最適化

from torch.optim import Adam

optimizer = Adam(params=model.parameters(), lr=2e-5)

スケジューラの定義

from transformers import get_linear_schedule_with_warmup

lr_scheduler = get_linear_schedule_with_warmup(
      optimizer=optimizer,

      num_warmup_steps=num_warmup_steps,

      num_training_steps=num_training_steps)

optimizer（obj）：AdamなどのPyTorchオプティマイザ
num_warmup_steps（int）：lrを線形増加させるステップ数。int(num_training_steps * 0.1)を設定
num_training_steps（int）：総学習ステップ数。len(train_dataloader) * num_epochsを設定

効率的な学習のためのモデル準備

prepareメソッドがデバイス配置を処理します

model, optimizer, dataloader, lr_scheduler = \
    accelerator.prepare(model,

                        optimizer,

                        dataloader,

                        lr_scheduler)

Acceleratorを使った学習ループの構築

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)

勾配をゼロに初期化
データをデバイスへ移動済み

Acceleratorを使った学習ループの構築

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

勾配をゼロに初期化
データをデバイスへ移動済み
データを手動で移動する行を削除

Acceleratorを使った学習ループの構築

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)

    loss = outputs.loss
    loss.backward()

勾配をゼロに初期化
データをデバイスへ移動済み
データを手動で移動する行を削除
順伝播の実行
クロスエントロピー損失と勾配の計算

Acceleratorを使った学習ループの構築

for batch in dataloader:

    optimizer.zero_grad()

    inputs, targets = batch

    outputs = model(inputs)
    loss = outputs.loss

    accelerator.backward(loss)

    optimizer.step()
    scheduler.step()

勾配をゼロに初期化
データをデバイスへ移動済み
データを手動で移動する行を削除
順伝播の実行
クロスエントロピー損失と勾配の計算
loss.backwardをacceleratorに置き換え
モデルパラメータと学習率を更新

変更点のまとめ

Accelerator導入前

データを手動でデバイスへ移動する必要あり
- inputs.to(device)
- targets.to(device)
loss.backward()で勾配を計算

Accelerator導入後

デバイス配置とデータ並列処理の自動化
- accelerator.prepare(model)
- accelerator.prepare(dataloader)
accelerator.backward(loss)で勾配の同期を処理
カスタマイズ可能なループ
使いやすく、ハードウェア非依存、スケーラブル、保守性が高い

では、練習しましょう！

PyTorch による効率的な AI モデルトレーニング