基本的な PyTorch による混合精度トレーニング

低精度の浮動小数点データ型を使用して、言語翻訳モデルのトレーニングを高速化しましょう。例えば、16ビット浮動小数点データ型（float16）は、32ビットの対応型（float32）の半分のサイズです。これにより、行列の乗算や畳み込みの計算が高速化されます。この処理では、勾配のスケーリングと演算の16ビット浮動小数点へのキャストが必要になることを思い出してください。

dataset、model、dataloader、optimizer は事前に読み込まれています。

この演習はコースの一部です

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

コースを見る

演習の手順

ループの前に、torch.amp.GradScaler を使用して勾配のスケーラーを定義します。
ループ内で、torch.autocast をコンテキストマネージャーとして使用し、演算を16ビット浮動小数点データ型にキャストします。
ループ内で、損失をスケーリングし、.backward() を呼び出してスケーリングされた勾配を生成します。

実践的なインタラクティブ演習

このサンプルコードを完成させて、この演習に挑戦してみましょう。

# Define a scaler for the gradients
scaler = torch.amp.____()
for batch in train_dataloader:
    inputs, targets = batch["input_ids"], batch["labels"]
    # Casts operations to mixed precision
    with torch.____(device_type="cpu", dtype=torch.____):
        outputs = model(inputs, labels=targets)
        loss = outputs.loss
    # Compute scaled gradients
    scaler.____(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

コードを編集して実行