PyTorchトレーニングループの解説

←エッセイ←エッセイ3クラスのスパイラルデータセット。陰影のある領域は、モデルのソフトマックス信頼度を示しています。トレーニングが進むにつれて、境界がより明確になります。PyTorchのトレーニングループを構築するのはかなり簡単ですが、すべての要素を正しい位置に正しい順序で配置することは、驚くほど不安定に感じられることがあります。多くの可動部分があり、最も基本的なエラーが修正された後でも、他のほとんどの間違いは発見するのが非常に難しい場合があります。コードの行が誤って配置されていると、トレーニングの実行が収束に失敗したり、誤った結果を生成したり、過剰なメモリを消費したりします。以下のセクションでは、各操作を順に説明し、各セクションの記述方法と、注意すべき一般的な間違いをすべて説明します。分散トレーニング、FSDP、マルチGPUセットアップはここでは範囲外ですが、将来のエッセイでこれらについては改めて説明します。（上記のアニメーションは、合成データでループを実行し、各エポックでの決定境界をキャプチャすることで生成されました。） 完全なループまず、完全なトレーニングループを見てみましょう。まだ理解したり記憶したりする必要はありません。ただ構造を感じ取ってください。 ```python 1import torch 2import torch.nn as nn 3from torch.utils.data import DataLoader, TensorDataset 4 5# --- data --- 6dataset = TensorDataset(X_train, y_train) 7loader = DataLoader(dataset, batch_size=64, shuffle=True) 8 9# --- model, loss, optimiser --- 10model = MLP(in_features=2, hidden=128, out_features=3) 11criterion = nn.CrossEntropyLoss() 12optimiser = torch.optim.Adam(model.parameters(), lr=1e-3) 13scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=100) 14 15# --- loop --- 16for epoch in range(100): 17 model.train() 18 for X_batch, y_batch in loader: 19 optimiser.zero_grad() 20 logits = model(X_batch) 21 loss = criterion(logits, y_batch) 22 loss.backward() 23 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) 24 optimiser.step() 25 scheduler.step() 26 27 model.eval() 28 with torch.no_grad(): 29 val_logits = model(X_val) 30 val_loss = criterion(val_logits, y_val) ``` それでは、各行が何をするのか、そしてどのように壊さないようにするのかを理解していきましょう。まず、いくつかの一般的な間違いから始めます。TL;DR 順序が本当に重要な場所最も一般的な失敗のいくつか、そして少し配置を間違えるだけでトレーニングループをどのように壊してしまうかを示します。これらを記憶する理由は、それらのどれも例外を発生させないためです。時間が経てば、トレーニング実行中にどのような種類のエラーを探すべきか分かるようになるでしょうが、最初のうちはこのカンニングシートが役立つはずです。 以下は最も一般的な失敗例と、配置を少し間違えることでトレーニングループがどのように壊れてしまうかを示します。 * `model.to(device)`を`optimiser = ...`の後に置く: dtype変換（例：`model.half()`）と組み合わせると、`nn.Module.to()`は新しい`nn.Parameter`オブジェクトを割り当てます。この場合、オプティマイザは破棄された元のオブジェクトへの参照を保持し、そちらに更新を適用してしまいます。 * `optimiser.zero_grad()`を`loss.backward()`の後に置く: 複数のバッチからの勾配が蓄積されます。更新は現在のバッチ単独ではなく、それらの合計を使用します。 * `clip_grad_norm_()`を`loss.backward()`の前に置く: `.grad`が空であるため、この呼び出しは何も行いません。 * `clip_grad_norm_()`を`optimiser.step()`の後に置く: すでに適用された勾配をクリップするため、効果がありません。 * `scheduler.step()`をバッチループの内側に置く: 学習率が1エポックに1回ではなく、`len(loader)`回減衰してしまいます。 * `model.eval()`の後に`model.train()`を省略する: Dropoutが無効になり、BatchNormがフリーズします。モデルはエラーなしで評価モードのままトレーニングされてしまいます。 * 検証中に`torch.no_grad()`を省略する: すべての検証バッチでAutogradグラフが構築されます。メモリがOOMになるまで増加します。 * `loss.item()`の代わりに`loss`をログに記録する: ロギング呼び出しの間、計算グラフがメモリに固定されてしまいます。 それでは、これらのそれぞれについて詳しく見ていきましょう。 データプラクティス PyTorchカスタムデータセットの簡単Q344 `__len__`と`__getitem__`を持つカスタムPyTorchデータセットを実装します。PyTorchのデータパイプラインには、DatasetとDataLoaderの2つの部分があります。Datasetは、`__len__`（データセット内の要素数）と`__getitem__`（予想通り、項目を取得します）を実装する単なるPythonオブジェクトです。それはテンソルをラップする単純なラッパーであることも、必要に応じてディスクからデータをロードすることもできます。DataLoaderはデータセットをラップし、バッチを生成します。データセット全体を一度通過するごとに1エポックとなります。`shuffle=True`の場合、各エポックで異なる順序でサンプルが提示されます。 ```python 1dataset = TensorDataset(X_train, y_train) 2loader = DataLoader( 3 dataset, 4 batch_size=64, 5 shuffle=True, 6 num_workers=2, 7 pin_memory=True, 8 persistent_workers=True, 9) ``` プラクティスPyTorch DataLoader 簡単Q342 PyTorch DataLoaderを設定する：バッチ処理、シャッフル、イテレーション。TensorDatasetは入力テンソルとラベルテンソルをインデックスでペアにします。`dataset[i]`でインデックスを指定すると`(X[i], y[i])`が返されます。DataLoaderは`__getitem__`を繰り返し呼び出し、結果をバッチにまとめ、オプションでバックグラウンドのワーカープロセスに処理を渡します。`num_workers`は、GPU計算と並行してバッチをプリフェッチする個別のプロセスを生成します。メインプロセスは、バッチがまだ準備できていない場合にのみ`.next()`でブロックします。ワーカー数がゼロの場合、メインプロセスがすべてのロードを実行し、データ量の多いタスクではGPU利用率のボトルネックになることがよくあります。2〜4個のワーカーが実用的ですが、適切な数はCPU数とI/O速度に依存します。`pin_memory=True`は、バッチテンソルをピン留めされたホストメモリに割り当てます。GPU DMAエンジンは、カーネルバッファを介して最初にコピーすることなく、ピン留めされたメモリから直接転送できるため、ホストからデバイスへの転送時間を短縮します。これは、`num_workers > 0`でCUDAに転送している場合にのみ役立ちます。`persistent_workers=True`は、エポック間でワーカープロセスをアクティブに保ちます。これがないと、各エポックの開始時にワーカーが再起動され、多数のワーカーを使用する際に測定可能なフォークオーバーヘッドが追加されます。`drop_last=True`は、最終バッチが`batch_size`よりも小さい場合にそのバッチを破棄します。2、3のサンプルバッチから計算されたBatchNorm統計はノイズが多いため、残りを破棄することでこれを回避します。データを破棄するという点でわずかなコストがかかりますが、安定性のためにはしばしば価値があります。小さなバッチはノイズの多い勾配推定値を生成し、暗黙の正則化として機能します。大きなバッチはより多くのGPUメモリを使用しますが、より多くの並列処理を可能にします。最も重要な効率の1つは、2の冪乗がテンソルコアのタイルサイズ（通常、dtypeに応じて16×16または8×16）と一致することを知っていることです。そのため、バッチサイズとレイヤーの次元を8または16の倍数に設定することは良いアイデアです。`.to(device)`はテンソルをターゲットデバイスに移動します。テンソルの場合、インプレースではありません。新しいテンソルを返し、元のテンソルは変更されません。たとえば、`X_batch.to('cuda')`はGPU上に新しいテンソルを返しますが、`X_batch`自体はCPU上に残ります。 再現性 モデルとローダーを構築する前にシードを設定すると、すべての実行で同じ結果が得られます。これは、実験を再現し、モデルの動作が決定論的であることを確認するために不可欠です。これがモデルに影響を与える主な領域は、データローダーとモデルの重みの初期化です。 ```python 1import random 2import numpy as np 3 4def set_seed(seed: int = 42): 5 torch.manual_seed(seed) 6 torch.cuda.manual_seed_all(seed) 7 np.random.seed(seed) 8 random.seed(seed) 9 torch.backends.cudnn.deterministic = True 10 torch.backends.cudnn.benchmark = False 11 12set_seed(42) ``` `torch.manual_seed`はCPUジェネレーターをシードします。`torch.cuda.manual_seed_all`はすべてのGPUをシードします。NumPyとPythonの`random`はPyTorchが触れない独立したRNGであるため、注意が必要です。それらには別の乱数シードが必要になる場合があります。`cudnn.deterministic = True`はcuDNNに決定論的な畳み込みアルゴリズムの使用を強制します。一部のcuDNNカーネルは、スループットのためにデフォルトで非決定論的です。決定論的な代替手段はわずかに遅いですが、開発中は実質的にそれほど問題にならないはずです。`cudnn.benchmark = False`は`deterministic = True`と組み合わせる必要があります。`benchmark = True`の場合、cuDNNは入力形状ごとにいくつかのアルゴリズムをプロファイリングし、最速のものを選択します。このプロセス自体は実行ごとに異なります。これを`False`に固定すると、常に同じ結果が得られることが保証されます。`num_workers > 0`の場合、各DataLoaderワーカーは、プロセスをフォークするときにOSによってシードされる独自のRNG状態を持ちます。ワーカーのランダム性を再現可能にするには、ジェネレーターと`worker_init_fn`を渡す必要があります。