PyTorch学習ループ チートシート

このシートは、PyTorchのループを書いたことがあり、操作の順序と危険な箇所を1ページで確認したい人のためのものです。すべてのブロックは合成データを使い、CPU上で数秒で動きます。ループの背景にある考え方は /learn/beginner/04-training-loop-and-loss/ と /learn/beginner/07-your-first-neural-network/ を参照してください。

定番のループ

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

torch.manual_seed(0)

# Synthetic regression data: y = 3 x0 - 2 x1 + noise
x = torch.randn(512, 2)
y = (3 * x[:, 0] - 2 * x[:, 1] + 0.1 * torch.randn(512)).unsqueeze(1)
train_loader = DataLoader(TensorDataset(x[:400], y[:400]), batch_size=32, shuffle=True)
val_loader = DataLoader(TensorDataset(x[400:], y[400:]), batch_size=64)

model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
loss_fn = nn.MSELoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-2, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

for epoch in range(20):
    model.train()                              # dropout and batch norm in training mode
    for xb, yb in train_loader:
        optimizer.zero_grad(set_to_none=True)  # clear last step's gradients
        pred = model(xb)                       # forward pass
        loss = loss_fn(pred, yb)               # scalar loss
        loss.backward()                        # gradients land in .grad
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()                       # update weights
    scheduler.step()                           # once per epoch, after optimizer.step()

    model.eval()                               # inference mode for layers that care
    total, n = 0.0, 0
    with torch.no_grad():                      # no graph, no .grad, less memory
        for xb, yb in val_loader:
            total += loss_fn(model(xb), yb).item() * len(xb)
            n += len(xb)
    if epoch % 5 == 4:
        print(f"epoch {epoch + 1:2d}  val_mse {total / n:.4f}  lr {scheduler.get_last_lr()[0]:.5f}")

epoch  5  val_mse 0.4046  lr 0.00854
epoch 10  val_mse 0.0395  lr 0.00500
epoch 15  val_mse 0.0286  lr 0.00146
epoch 20  val_mse 0.0270  lr 0.00000

1行ずつ、重要な点を見ていきます。

行	なぜ必要か	ないと何が起きるか
model.train()	ドロップアウトを有効にし、バッチ正規化をバッチ統計量に切り替える	検証データが移動統計量に漏れ込む、あるいはドロップアウトが無効のままで過学習する
optimizer.zero_grad(set_to_none=True)	勾配は設計上、backward()の呼び出しをまたいで累積する	毎ステップが過去のすべての勾配の合計を使う。損失が発散するか停滞する
loss.backward()	autogradグラフをたどり、すべての葉パラメータの.gradを埋める	更新に使うものがない。同じグラフで2回呼ぶと、retain_graph=Trueでない限りエラーになる
clip_grad_norm_	全勾配のグローバルノルムが最大max_normになるように再スケールする	悪いバッチによるまれな大きい勾配が重みを壊す。backward()の後、step()の前にクリップする
optimizer.step()	.gradを持つすべてのパラメータに更新則を適用する	重みがまったく変わらない
scheduler.step()	学習率をスケジュールに沿って進める	optimizer.step()より前に呼ぶと最初の値が飛ばされ、警告が出る。ステップ単位のスケジューラ（OneCycleLR）はバッチループの中に置く
model.eval()	ドロップアウトを無効にし、バッチ正規化は移動統計量を使う	検証の数値がノイズだらけで悲観的すぎる値になる
torch.no_grad()	autogradの記録を止める	検証が解放されないグラフを作り続け、メモリが増え、速度も落ちる
.item()	デバイスからPythonのfloatを取り出す	テンソルのまま持つとそのグラフも生き続ける。エポックをまたいで損失テンソルを合計するのは典型的なリーク

torch.inference_mode()は、より厳格で少し速いno_grad()です。その中で作られたテンソルは後でautogradに使えません。評価には問題ありませんが、逆伝播させるものには不適切です。

混合精度

autocastは行列積と畳み込みを低精度で実行し、リダクション、softmax、マスターの重みはfp32のままにします。CPUで対応している低精度はbf16なので、このブロックはどこでも動きます。CUDAではdevice_typeを変えて、AmpereかそれよりあとのGPUならbf16のままにします。fp16を使わざるを得ない場合（古いGPU）はtorch.amp.GradScalerを追加し、素の呼び出しの代わりにscaler.scale(loss).backward()、scaler.step(optimizer)、scaler.update()を呼び、クリッピングの前にscaler.unscale_(optimizer)でスケールを戻します。

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(64, 128), nn.GELU(), nn.Linear(128, 10))
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
x = torch.randn(32, 64)
target = torch.randint(0, 10, (32,))

# bf16 autocast works on CPU, so the same code runs on a laptop and a GPU.
# On CUDA use device_type="cuda"; for float16 add a torch.amp.GradScaler.
device_type = "cpu"
for step in range(3):
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
        logits = model(x)                       # matmuls run in bf16
        loss = nn.functional.cross_entropy(logits, target)  # reductions stay fp32
    loss.backward()                             # outside autocast
    optimizer.step()
    print(f"step {step}  loss {loss.item():.4f}  logits dtype {logits.dtype}  weights {model[0].weight.dtype}")

出力にはlogits dtype torch.bfloat16とweights torch.float32が表示されます。活性はキャストされ、パラメータはされません。backward()とstep()は意図的にコンテキストの外に置いています。

チェックポイント

重みだけでなく、再開に必要なものをすべて保存します。信頼できないソースのチェックポイントがコードを実行できないように、weights_only=True（PyTorch 2.6以降は既定）で読み込みます。

import tempfile, pathlib
import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(4, 2)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)

# one fake step so the optimiser has state worth saving
loss = model(torch.randn(8, 4)).pow(2).mean()
loss.backward()
optimizer.step()
scheduler.step()

ckpt = {
    "epoch": 1,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "rng": torch.get_rng_state(),
}
path = pathlib.Path(tempfile.mkdtemp()) / "ckpt.pt"
torch.save(ckpt, path)

restored = nn.Linear(4, 2)
state = torch.load(path, map_location="cpu", weights_only=True)
restored.load_state_dict(state["model"])
same = all(torch.equal(a, b) for a, b in zip(model.state_dict().values(), restored.state_dict().values()))
print("epoch", state["epoch"], "weights identical:", same, "optimizer keys:", list(state["optimizer"].keys()))

epoch 1 weights identical: True optimizer keys: ['state', 'param_groups']と表示されます。DataParallelやDistributedDataParallelでラップしたモデルは、キーにmodule.接頭辞が付かないようにmodel.module.state_dict()を保存します。再開ではなく重みの共有が目的なら、safetensorsを使ってください（/reference/frameworks/ を参照）。

再現性

import random
import numpy as np
import torch
from torch import nn

def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)            # also seeds CUDA and MPS when present
    torch.use_deterministic_algorithms(True, warn_only=True)

def run(seed: int) -> float:
    seed_everything(seed)
    g = torch.Generator().manual_seed(seed)   # a generator for the DataLoader shuffle
    x = torch.randn(64, 8, generator=g)
    model = nn.Linear(8, 1)
    return model(x).sum().item()

print("same seed, same result:", run(1) == run(1))
print("different seed, different result:", run(1) != run(2))

どちらの行もTrueを表示します。CUDAでは、決定的な行列積のために環境変数CUBLAS_WORKSPACE_CONFIG=:4096:8とtorch.backends.cudnn.benchmark = Falseも必要です。GPU機種やドライババージョンをまたいでビット単位で一致させるのは現実的な目標ではありません。同じマシン、同じソフトウェア、同じシードでの一致が妥当な目標です。

DataLoaderの落とし穴

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

torch.manual_seed(0)
ds = TensorDataset(torch.arange(10).float().unsqueeze(1), torch.arange(10))

# drop_last=False leaves a short final batch, which skews a mean of per-batch losses
sizes = [len(xb) for xb, _ in DataLoader(ds, batch_size=4)]
print("batch sizes without drop_last:", sizes)

# a generator makes the shuffle order reproducible without touching the global seed
g = torch.Generator().manual_seed(123)
order = [int(yb[0]) for _, yb in DataLoader(ds, batch_size=1, shuffle=True, generator=g)]
print("shuffled order:", order)

# num_workers=0 on CPU demos; raise it only when the dataset does real decoding work
loader = DataLoader(ds, batch_size=5, num_workers=0, pin_memory=False)
print("batches per epoch:", len(loader))

batch sizes without drop_last: [4, 4, 2]
shuffled order: [2, 1, 6, 0, 4, 7, 9, 8, 3, 5]
batches per epoch: 2

• 損失の平均はバッチ数ではなくサンプル数で取ります。または学習時はdrop_last=Trueにします。
• num_workersが0より大きいとプロセスをフォークします。macOSとWindowsではデータセットとcollate関数がimport可能でなければならず（ラムダ不可）、学習スクリプトにはif __name__ == "__main__":のガードが必要です。
• 各ワーカーは独自のnumpy状態を持ちます。worker_init_fnでシードを設定しないと、ランダムなデータ拡張がワーカー間で同じになります。
• pin_memory=TrueはCUDAへのコピー時にしか効果がありません。persistent_workers=Trueはエポックごとのワーカー再生成を避けます。
• 検証ローダーのshuffle=Trueは害はありませんが無意味です。学習ローダーのshuffle=Falseはバグです。

よくあるエラー

エラーまたは症状	原因	対処
RuntimeError: Expected all tensors to be on the same device	モデルがGPU、バッチがCPU（またはその逆）	両方を移す: model.to(device)を1回、xb.to(device)をバッチごと
Trying to backward through the graph a second time	前のステップのテンソルを損失で再利用している（バッチをまたいで保持した隠れ状態が典型）	持ち越すものを.detach()する
数ステップ後に損失がnanになる	学習率が高すぎる、fp16のオーバーフロー、ゼロの対数	学習率を下げる、bf16かGradScalerを使う、対数の中にイプシロンを足す
損失がln(num_classes)で止まる	ラベルのdtypeや形状が違う、またはzero_gradがない	ラベルは形状(N,)のint64クラスID、ロジットは(N, C)
学習損失も検証損失も下がるのにテストが悪い	分割でのリーク	/learn/beginner/02-data-is-the-model/ を参照
検証でCUDA out of memory	検証がno_grad()の外にある、またはバッチが学習より大きい	torch.no_grad()で包む、バッチサイズを揃える
エポックごとにメモリが増える	.item()のfloatではなく損失テンソルを蓄積している	.item()か.detach()を呼ぶ
UserWarning: Detected call of lr_scheduler.step() before optimizer.step()	順序が逆、またはGradScalerが最初のステップをスキップした	オプティマイザを先にステップする。スケーラーがステップをスキップした場合の警告は無害
実行ごとに結果が違う	シードなし、非決定的な演算、ワーカーのシード	上の再現性のブロック
view size is not compatible	permuteの後でテンソルが非連続になっている	.contiguous()か.reshape()。/reference/tensor-shapes-and-broadcasting/ を参照

出典

• torch.optimドキュメント（スケジューラの順序を含む）
• 自動混合精度パッケージ torch.amp
• モデルの保存と読み込み
• 再現性
• torch.utils.dataドキュメント
• Micikevicius et al., Mixed Precision Training, 2017
• Loshchilov and Hutter, Decoupled Weight Decay Regularization (AdamW), 2017