AMPによる混合精度学習
torch.cuda.amp.autocast()、GradScaler、FP16とBF16の違い、メモリ削減、速度向上について学習します。
「AMPによる混合精度学習」はCoddyKit上の無料Learn AI with Pythonレッスンです。 これはレッスン3/4です。 下記で完全なレッスンを無料で読むことができます。その後、ブラウザ内の組み込みコードエディタと24時間対応のAIチューターでハンズオン演習できます。 これはLearn AI with Python学習パスの一部であり、ウェブとCoddyKitアプリ全体で進捗が同期されます。 Learn AI with Pythonコースには全4レッスンが含まれています。
混合精度とは
混合精度学習では、ほとんどの演算を32ビット浮動小数点（FP32）ではなく16ビット浮動小数点（FP16）で実行します。FP16はメモリ使用量が半分になり、最新のGPUのTensor Coreでより高速に実行できます。一方、数値的に敏感な一部の演算は、安定性を保つためFP32のまま実行します。
メリット
混合精度には次のメリットがあります。
- メモリを約2倍節約でき、より大きなバッチやモデルを使用できます
- Tensor Coreで行列乗算を高速化できます
- 正しく使用すれば、最終的なモデル精度の低下はほとんど、またはまったくありません
FP16のアンダーフロー問題
FP16の値域は狭いという特徴があります。小さな勾配値がアンダーフローしてゼロになると、学習が気付かないうちに停止します。これが混合精度で解決すべき中心的な課題であり、すべてを単純にFP16へ変換できない理由です。
torch.cuda.amp
PyTorchのAutomatic Mixed Precision（AMP）は、2つのツールでこの処理を行います。autocastは演算ごとに精度を選択し、GradScalerは勾配のアンダーフローを防ぎます。
import torch
from torch.cuda.amp import autocast, GradScalerautocastコンテキスト
フォワードパスをautocast()で囲みます。その内部では、PyTorchが各演算を最も安全な精度で自動的に実行します。行列乗算はFP16で、softmaxや損失のような集約演算はFP32で実行されます。
with autocast():
output = model(x)
loss = criterion(output, y)GradScalerの導入
GradScalerは、バックプロパゲーションの前に損失へ大きなスケール係数を掛けることで、アンダーフローを防ぎます。これにより、小さな勾配をFP16で表現可能な範囲へ移し、オプティマイザのステップ前にスケールを元に戻します。
scaler = GradScaler()損失のスケーリング
scaler.scale(loss).backward()は損失を大きくしてから、バックプロパゲーションを実行します。結果として得られる勾配もスケーリングされるため、小さな値が消失するのを防げます。
scaler.scale(loss).backward()scaler.step
scaler.step(optimizer)は、まず勾配を本来の大きさに戻し、その後optimizer.step()を呼び出します。ただし、infやNaNが発生していない場合に限ります。勾配がオーバーフローした場合は、ステップをスキップします。
scaler.step(optimizer)scaler.update
scaler.update()は、次の反復に向けてスケール係数を調整します。ステップが成功するとスケールを大きくし、オーバーフローが発生すると小さくします。この適応的なスケーリングにより、学習の安定性が自動的に保たれます。
scaler.update()AMPを使った完全なループ
これらの要素を組み合わせると、混合精度学習の完全なステップになります。
scaler = GradScaler()
for x, y in loader:
x, y = x.cuda(), y.cuda()
optimizer.zero_grad()
with autocast():
out = model(x)
loss = criterion(out, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()実践的なヒント
次の点を覚えておきましょう。
- スケーリングした損失に対してのみ
backward()を呼び出します - autocastで囲むのはフォワードパスだけです。バックワードやオプティマイザのステップは囲みません
- AMPはTensor Core搭載GPUで特に効果を発揮します。古いハードウェアでは効果が小さくなります
確認テスト
AMPの理解度を確認しましょう。
まとめ
AMPを使った混合精度学習について学びました。
autocast()はフォワードパスで、演算ごとにFP16またはFP32を選択します- GradScalerは損失をスケーリングして、勾配のアンダーフローを防ぎます
- 一連の処理は
scaler.scale(loss).backward()、scaler.step(optimizer)、scaler.update()です - 結果として、メモリを約2倍節約でき、学習も高速になります
よくある質問
「AMPによる混合精度学習」レッスンは無料ですか？
はい。「AMPによる混合精度学習」の完全なテキストはこのウェブで無料で読めます。インタラクティブに演習し（組み込みコードエディタと24時間対応のAIチューター）、Learn AI with Pythonコースの残りをアンロックするには、CoddyKit PROにアップグレードしてください。 Learn AI with Pythonコースには全4レッスンが含まれています。
「AMPによる混合精度学習」で何を学びますか？
torch.cuda.amp.autocast()、GradScaler、FP16とBF16の違い、メモリ削減、速度向上について学習します。 ブラウザで直接実行するハンズオンコードでLearn AI with Pythonを演習し、24時間対応のAIチューターがレッスンを進める中での質問に答えます。
Learn AI with Pythonを始めるのに経験は必要ですか？
事前経験は必要ありません。CoddyKitのLearn AI with Pythonは初級者から上級者向けに構成されているため、ここから始めるか最初から始めて、自分のペースで進むことができます。 これはレッスン3/4です。
「AMPによる混合精度学習」レッスンにはどのくらい時間がかかりますか？
ほとんどのCoddyKitレッスンは約5～10分かかります。各レッスンはコンパクトでインタラクティブなので、着実に進歩し、ウェブとアプリ全体で正確に前回の場所から再開できます。
このLearn AI with Pythonレッスンでコードを書いて実行できますか？
はい。すべてのLearn AI with Pythonレッスンに組み込みコードエディタが含まれているため、ブラウザでリアルコードを書いて実行し、即座のAIフィードバックを取得できます。ローカル設定は不要です。