AdamWによるバランスの取れた学習

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

Dennis Lee

Data Engineer, Amazon

効率的な学習

オプティマイザに焦点を当てたコースの章トピックを示す図。

学習効率のためのオプティマイザ

AdamW、Adafactor、8-bit Adamの3つのオプティマイザを示す図。

学習効率のためのオプティマイザ

AdamW、Adafactor、8-bit Adamの3つのオプティマイザを示す図。

学習効率のためのオプティマイザ

AdamW、Adafactor、8-bit Adamの3つのオプティマイザを示す図。

オプティマイザのトレードオフ

AdamW、Adafactor、8-bit AdamのパラメータD数と精度のトレードオフを示す図。

オプティマイザのトレードオフ

AdamW、Adafactor、8-bit AdamのパラメータD数と精度のトレードオフを示す図。

オプティマイザのトレードオフ

AdamW、Adafactor、8-bit AdamのパラメータD数と精度のトレードオフを示す図。

オプティマイザのトレードオフ

AdamW、Adafactor、8-bit AdamのパラメータD数と精度のトレードオフを示す図。

AdamWの仕組み

AdamWの仕組みを示す図。

AdamWの仕組み

AdamWの仕組みを示す図。

AdamWの仕組み

AdamWの仕組みを示す図。

勾配の指数移動平均（EMA）を計算する

AdamWの仕組み

AdamWの仕組みを示す図。

勾配の指数移動平均（EMA）を計算する
勾配の二乗のEMAを計算する

AdamWの仕組み

AdamWの仕組みを示す図。

勾配の指数移動平均（EMA）を計算する
勾配の二乗のEMAを計算する

AdamWの仕組み

AdamWの仕組みを示す図。

勾配の指数移動平均（EMA）を計算する
勾配の二乗のEMAを計算する

AdamWのメモリ使用量

AdamWの計算に関わるパラメータ勾配・勾配のEMA・勾配の二乗のEMAを示す図。

各マスはパラメータ、各色は状態を表す

AdamWのメモリ使用量

AdamWの計算に関わるパラメータ勾配・勾配のEMA・勾配の二乗のEMAを示す図。

各マスはパラメータ、各色は状態を表す

AdamWのメモリ使用量

AdamWの計算に関わるパラメータ勾配・勾配のEMA・勾配の二乗のEMAを示す図。

各マスはパラメータ、各色は状態を表す
パラメータあたりのメモリ = 8バイト = 4バイト/状態 × 2状態
総メモリ = パラメータあたりのメモリ（8バイト）× パラメータ数

AdamWのメモリ使用量を推定する

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

num_parameters = sum(p.numel() for p in model.parameters())
print(f"Number of model parameters: {num_parameters:,}")

Number of model parameters: 65,783,042

estimated_memory = num_parameters * 8 / (1024 ** 2)
print(f"Estimated memory usage of AdamW: {estimated_memory:.0f} MB")

Estimated memory usage of AdamW: 502 MB

TrainerとAccelerator

AcceleratorとTrainerのカスタマイズ性と使いやすさのトレードオフを示す図。

TrainerでAdamWを実装する

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


trainer = Trainer(model=model,
                  args=training_args,
                  train_dataset=train_dataset,
                  eval_dataset=validation_dataset,
                  compute_metrics=compute_metrics,
                  optimizers=(optimizer, lr_scheduler))

trainer.train()

{'epoch': 1.0, 'eval_accuracy': 0.7, 'eval_f1': 0.8}

AcceleratorでAdamWを実装する

from torch.optim import AdamW

optimizer = AdamW(params=model.parameters())


for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    outputs = model(inputs, labels=targets)
    loss = outputs.loss
    accelerator.backward(loss)
    optimizer.step()
    lr_scheduler.step()
    optimizer.zero_grad()
    print(f"Loss = {loss}")

Loss = 0.7

オプティマイザの状態を確認する

optimizer_state = optimizer.state.values()
print(optimizer_state)

dict_values([{'step': tensor(3.),

              'exp_avg': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...]),

              'exp_avg_sq': tensor([[0., 0., 0.,  ..., 0., 0., 0.], ...])}, ...])

オプティマイザのサイズを計算する

def compute_optimizer_size(optimizer_state):
    total_size_megabytes, total_num_elements = 0, 0

    for params in optimizer_state:

        for name, tensor in params.items():
            tensor = torch.tensor(tensor)

            num_elements = tensor.numel()

            element_size = tensor.element_size()

            total_num_elements += num_elements

            total_size_megabytes += num_elements * element_size / (1024 ** 2)

    return total_size_megabytes, total_num_elements

オプティマイザのサイズを計算する

total_size_megabytes, total_num_elements = \
    compute_optimizer_size(trainer.optimizer.state.values())
print(f"Number of optimizer parameters: {total_num_elements:,}")

Number of optimizer parameters: 131,566,188

print(f"Optimizer size: {total_size_megabytes:.0f} MB")

Optimizer size: 502 MB

さあ、練習しましょう！

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