PyTorch トレーニングループ完全ガイド：三クラススパイラルデータセット向け

記事概要

本チュートリアルでは、三クラススパイラルデータセットを用いたトレーニングの仕組みを解説します。暗黙的な陰影部分はモデルのソフトマックス信頼度を示しており、学習が進むにつれて境界線がシャープになる様子を可視化できます。

PyTorch のトレーニングループを作成するのは比較的容易ですが、要素を適切な場所と順序で配置することは驚くほど壊れやすいものです。以下に示す構成に従うことで、エラーの未然防止や学習効率の最大化が可能になります。

完全なループ構造

まず、トレーニング全体の骨格となるループ構造を確認します。理解や暗記が必須ではありませんので、全体の構造に対する感覚をつかんでおいてください。

1 import torch
2 import torch.nn as nn
3 from torch.utils.data import DataLoader, TensorDataset

4 # --- データ処理 ---
5 dataset = TensorDataset(X_train, y_train)
6 loader  = DataLoader(dataset, batch_size=64, shuffle=True)

7 # --- モデル、損失関数、最適化器 ---
8 model     = MLP(in_features=2, hidden=128, out_features=3)
9 criterion = nn.CrossEntropyLoss()
10 optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)
11 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=100)

12 # --- トレーニングループ ---
13 for epoch in range(100):
14     model.train()
15     for X_batch, y_batch in loader:
16         optimiser.zero_grad()
17         logits = model(X_batch)
18         loss   = criterion(logits, y_batch)
19         loss.backward()
20         torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
21         optimiser.step()
22     scheduler.step()

23     model.eval()
24     with torch.no_grad():
25         val_logits = model(X_val)
26         val_loss   = criterion(val_logits, y_val)

注意点：順序が極めて重要です

ここにはトレーニングループを崩す最も一般的な失敗例があります。配置を少し誤るだけで、例外エラーだけでなく以下のような問題を引き起こします。

grad

len(loader)

1. データ処理

PyTorch のデータパイプラインには

Dataset

DataLoader

• Dataset: 単なる Python オブジェクトで、

  __len__

  と

  __getitem__

  を実装します。
• DataLoader: データセットをラップし、バッチ生成、シャッフル、並列プリフェッチを行います。

PyTorch DataLoader の設定要素

loader = DataLoader(
    dataset,
    batch_size=64,             # バッチサイズ
    shuffle=True,               # 各エポックで順序をランダム化
    num_workers=2,              # プリフェッチ用のワーカープロセス数
    pin_memory=True,            # ホストメモリへのピン留め（非同期転送のため）
    persistent_workers=True     # エポック間でのワーカーの存続（再起動コスト削減）
)

重要なパラメータ解説

• TensorDataset

  : インデックスごとに入出力テンソルをペアリングします。

  dataset[i]

  は

  (X[i], y[i])

  を返します。

• num_workers

  : データの読み込みは GPU 計算と並行で行われます。ワーカー数がゼロの場合、メインプロセスが全ての読み込みを行うため、大量データではボトルネックになります。2 つ〜4 つが実用的です。

• pin_memory=True

  : バッチをピン留めされたホストメモリに割り当てます。これにより GPU DMA エンジンはピン留めメモリから直接転送でき、ホスト→デバイス転送時間が短縮されます（

  num_workers > 0

  の場合のみ有効）。

• persistent_workers=True

  : エポック間のワーカープロセスを存続させます。これにより各エポック開始時のフォークオーバーヘッドを防げます。

• drop_last=True

  : 最後のバッチが

  batch_size

  より小さい場合、それを棄却します。バッチ正規化の統計量にノイズを持たせないため、安定性の高い設定です。

バッチサイズと効率性について

• 小さすぎる: グラジエント推定量がノイズを増大させます（ただし暗黙的な正則化として機能）。
• 大きすぎる: GPU メモリ消費が増加しますが、並行処理を可能にします。
• 推奨: テンソルコアのタイルサイズ（通常 16×16 など）とバッチサイズ・レイヤー次元を8 または 16 の倍数にすると効率が良くなります。

再現性（シード設定）

実験の再現性を確保するために、モデル構築前にシードを設定してください。

def set_seed(seed: int = 42):
    torch.manual_seed(seed)              # CPU 生成子
    torch.cuda.manual_seed_all(seed)     # GPU 全機器
    import numpy as np, random
    np.random.seed(seed)                 # NumPy RNG
    random.seed(seed)                    # Python RNG
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

多プロセス環境での再現性:

num_workers > 0

worker_init_fn

2. モデル構築

nn.Module

• super().__init__()

  : 必須です。サブモジュールやパラメータテンソルを登録するためです。

• .to(device)

  : 最適化器作成前に呼び出してください。これで全てのパラメータがデバイスに移動します。
  • ※データ型変換（例：

    model.half().to(device)

    ）は内部で新しい

    nn.Parameter

    を生成するため、以前作成した最適化器の参照は無効になります。注意が必要です。

モデル構造とフック

• .parameters()

  : 学習パラメータ (

  requires_grad=True

  ) の返却に用います。

• forward(x)

  : デフォルトでは

  __call__

  が呼ばれ、フォワードフークを実行します。フック（計測やグラジエント修正など）を掛ける際は重要です。

• torch.compile()

  (PyTorch 2.0+): フォワードパスのトレースを行い、最適化されたカーネルを生成します。初期実行は遅くなりますが、以後 10〜30% の高速化が可能です。

3. トレーニングモード (

model.train()

)

モデル全体に

training=True

• Dropout: ランダムマスクを適用し、残存アクティベーションをスケールします（推論時ではアイデンティティ関数）。
• BatchNorm: バッチ統計量を計算・更新するか、蓄積された推定値 (

  running_mean/

  running_var

  ) を使用するかを決めます。

注意: 検証モードで

torch.no_grad()

を省略すると、メモリ使用量が激増します（自動グラジエント計算のため）。両方を併用するのが標準的です。

4. フォワードパスとグラフ構築

requires_grad=True

• 動的グラフ: 実行時にオンザフライで構築されるため、条件分岐や可変サイズモデルに対応します。
• メモリコスト: バッチサイズ

  N

  と深さ

  L

  に依存し、

  O(NL)

  でスケールします。
  • グラジエントチェックポイントング (

    checkpoint

    ): 中間アクティベーションを保存せず、必要に応じて再計算することでメモリの削減を行います（計算量の増加は発生しますが）。

from torch.utils.checkpoint import checkpoint
# ...
x = checkpoint(self.block1, x, use_reentrant=False)
# ...

5. 損失関数 (

CrossEntropyLoss

)

4 つのステップ目でスカラー損失を計算します。

CrossEntropyLoss

LogSoftmax

NLLLoss

• 主な引数:

  • weight

    : クラスの重み（不均衡データ用）。

  • label_smoothing

    : 過信を減らすための平滑化パラメータ。

  • ignore_index

    : パディングトークンを除外するための指定。

  • reduction='mean'

    : バッチ平均を行います（

    sum

    と切り替えると実質的な学習率が変わります）。

重要: ロギング時は必ず

loss.item()

6. バックワードパス（逆伝播）

loss.backward()

.grad

• 加法の累積: PyTorch はグラジエントを置換せず加算します。

• zero_grad()

  の必要性: 次のフォワードパス前に必ず

  .grad

  をリセット（または無効化）する必要があります。

  • zero_grad(set_to_none=True)

    : メモリ解放が早くなり、速度向上に寄与します（PyTorch 2.0 以降のデフォルト）。

思考:

retain_graph=True

は、同一グラフを複数回使用する場合（例：GAN）のみ必要です。通常は即座に解放されます。

7. グラジエントクリッピング

グラジエントスパイクを防ぐため、グローバル L2 ノルムを制限します。

• clip_grad_norm_(max_norm=1.0)

  : グラジエントベクトルの方向は維持しつつ、大きさを

  max_norm

  に制約します。
  • この処理は

    backward()

    後、

    step()

    前に行う必要があります。

8. 重みの更新 (

optimiser.step()

)

最適化器は

.grad

$$ \theta_t = \theta_{t-1} - \eta \cdot \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} $$

• Adam vs AdamW:
  • Adam: L2 正則化はグラジエント計算に含まれます。
  • AdamW: 重み減衰（L2 正則化）をグラジエント更新から分離し、直接重みに適用します。現在ではTransfomer モデルの標準実装です。

ハイパーパラメータと設定

• fused=True

  : CUDA カーネルへ融合して速度向上（約 30〜50%）。

• foreach=True

  : バッチ処理による CPU での計算加速。
• メモリ効率化: 巨大パラメータ数モデルでは、8 ビット浮動小数点オプティマイザ（bitsandbytes）や ZeRO を検討してください。

9. 学習率スケジュール

scheduler.step()

• 誤った用法: バッチループ内で呼び出すと、LR が過剰に減衰します。
• 推奨手法: リニアウォームアップ + コサイン退避（Cosine Annealing with Warmup）。
• Plateau Reducer:

  ReduceLROnPlateau

  は検証性能が一定期間改善しない場合に LR を低下させます（メトリックを引数に渡す必要あり）。

10. 検証ステップ (

model.eval()

)

学習ループ終了後に、モデルを検証モードに切り替えます。

model.eval()
with torch.no_grad():
    val_logits = model(X_val)
    val_loss   = criterion(val_logits, y_val).item()
    val_acc    = (val_logits.argmax(1) == y_val).float().mean()

• model.eval()

  : ドロップアウト無効化、バッチ正規化の蓄積推定値への切り替え。

• torch.no_grad()

  または

  torch.inference_mode()

  : グラフ構築停止によりメモリ使用量を約半分削減。推論モードの方が

  no_grad()

  よりもさらに高速です。

11. チェックポイント保存と復元

トレーニングの中断や再開を可能にします。

# セーブ
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimiser_state_dict': optimiser.state_dict(), # モメント状態も保存が必須
    'scheduler_state_dict': scheduler.state_dict(),
    'val_loss': val_loss,
}, 'checkpoint.pt')

# リロード
checkpoint = torch.load('checkpoint.pt', map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
optimiser.load_state_dict(checkpoint['optimiser_state_dict'])

重要: 最適化器の状態（モーメント推定量）も保存して復元しないと、トレーニングを正しい状態から再開できません（通常は損失スパイクを引き起こします）。ベストなモデルは検証性能が最も良かったエポックの重みをロードしてください。

12. GPU 効率性のテクニック

いくつかの基本テクニックを導入することで、ハードウェアの全性能を引き出せます。

デバイス配置

モデルとデータを同じデバイスに置き、転送オーバーヘッドを削減します。

non_blocking=True

ミクシドプレシジョン (Mixed Precision)

FP16 または BF16 を使用することで高速化を図ります。ただし動的範囲の問題から、GradScaler の活用が重要です。

scaler = torch.amp.GradScaler('cuda')

for X_batch, y_batch in loader:
    optimiser.zero_grad()
    
    with torch.amp.autocast('cuda', dtype=torch.float16):
        logits = model(X_batch)
        loss   = criterion(logits, y_batch)

    scaler.scale(loss).backward()     # スケーリング適用
    scaler.unscale_(optimiser)        # クリップ前解除
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimiser)            # ステップ（inf/nan時はスキップ）
    scaler.update()                   # スケール因子の調整

DataLoader パイプライン

num_workers

pin_memory

persistent_workers

コンパイル (

torch.compile

)

PyTorch 2.0+ において、トリートされたモデルを生成し、メモリアクセスを削減しながら操作融合を行うことで、10〜30% の速度向上が期待できます。

完全な注釈付きスクリプト（統合例）

以下のコードは、前述のセクションすべてを組み合わせた最終的なトレーニングループです。

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

# 1. データ処理 (TensorDataset -> DataLoader)
X_train = torch.randn(2000, 2)                                               # [1] 入力
y_train = make_labels(X_train)                                               # [2] ラベル
dataset = TensorDataset(X_train, y_train)                                    # [3] パーリング
loader  = DataLoader(dataset, batch_size=64, shuffle=True,                   # [4] バッチ生成
                     num_workers=2, pin_memory=True)

# 2. モデル構築 (__.init__ -> .to(device))
class MLP(nn.Module):                                                       # [5] 定義
    def __init__(self, in_features, hidden, out_features):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden),      nn.ReLU(),
            nn.Linear(hidden, out_features)
        )
    def forward(self, x): return self.net(x)

model = MLP(...).to(device)                                                 # [6] デバイス移動 (最適化器作成前)

# 3. 損失、最適化器、スケジューラー
criterion = nn.CrossEntropyLoss()                                           # [7] 損失関数
optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)                   # [8] パラメータ参照
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=100)

# 4. トレーニングループ (train -> batch -> step)
for epoch in range(NUM_EPOCHS):                                            # [9] エポック制御

    model.train()                                                           # [10] トレーニングモード
    
    for X_batch, y_batch in loader:                                         # [11] バッチ読み込み
        X_batch = X_batch.to(device, non_blocking=True)
        y_batch = y_batch.to(device, non_blocking=True)

        optimiser.zero_grad()                                               # [12] グラジエントクリア
        
        logits = model(X_batch)                                             # [13] フォワード (グラフ構築)
        
        loss = criterion(logits, y_batch)                                   # [14] 損失計算
        
        loss.backward()                                                     # [15] バックワーズ（逆伝播）
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)    # [16] グラジエントクリップ
        
        optimiser.step()                                                    # [17] 重み更新
        
    scheduler.step()                                                        # [18] 学習率調整（エポック終了）

# 5. 検証ループ (eval -> no_grad)
model.eval()                                                               # [19a] 評価モード
with torch.no_grad():                                                      # [19b] グラフ無効化
    val_logits = model(X_val.to(device))
    val_loss   = criterion(val_logits, y_val).item()

このように整理して記述することで、崩れやすいトレーニングループの安定性を保証しつつ、GPU 効率も最大化できます。