実験日: 2026-08-24 / PyTorch 2.13.0+cu130
GPU: RTX 5070 Ti 16 GB / WSL2 Ubuntu 26.04
ニューラルネットの学習では、データを1枚ずつではなくbatchにまとめます。ではbatch sizeは大きいほど速く、VRAMに入る最大値へ設定すればよいのでしょうか。
VRAMに余裕があっても、空き容量を埋め切るbatchが最適とは限りません。
FashionMNIST全60,000枚を使い、batch 32～8192でCNNを各3回学習しました。さらに合成batchを131,072まで増やして、物理VRAMを超えたときの性能低下と、実際のtorch.OutOfMemoryErrorを記録しました。
結論は、batch 256ですでに最大throughputの約93%へ達し、batch 1024～8192はほぼ横ばいでした。batch 8192は256より約6.5%速いだけなのに、peak allocatedは106 MiBから2,704 MiBへ約25倍に増えました。
どこで頭打ちになったか
batch 32
16,062 images/s / 1 epoch 3.735秒 / step 1.99 ms / allocated 70.7 MiB
batch 256
34,738 images/s / 1.727秒 / step 7.35 ms / allocated 106.2 MiB
batch 1024
37,124 images/s / 1.616秒 / step 27.39 ms / allocated 399.3 MiB
batch 2048（今回の最高中央値）
37,332 images/s / 1.607秒 / step 53.57 ms / allocated 728.9 MiB
batch 8192
36,983 images/s / 1.622秒 / step 202.80 ms / allocated 2,704.3 MiB
- batch 32→256ではthroughputが約2.16倍になりました。
- 256→2048は約7.5%しか増えません。
- 256は最大値の約93.1%を、allocated 106 MiBで得られました。
- 8192は256より約6.5%高速ですが、allocatedは約25.5倍です。
- WSL2では自然条件でPyTorch
reservedが物理VRAMを超えても処理が続き、batch 32,768→49,152で約19.9万→5.7万 images/sへ急落しました。 - 可視メモリの25%というallocator上限を明示すると、batch 16,384は成功し、24,576で実際のCUDA OOMになりました。
batch sizeを大きくすると何が変わるのか
batch sizeをBとすると、1回のforward/backwardでB個の入力を並列処理します。batchを大きくする主な利点は次の2つです。
- GPUへ渡す1回の仕事を大きくし、kernel launchなどの固定費を相対的に小さくする
- 大きな行列・畳み込みとして計算し、GPUの演算器やmemory bandwidthを使いやすくする
一方で、activationは入力ごとに保存するため、おおむねbatch sizeに応じてVRAM消費が増えます。また1 epochあたりのoptimizer更新回数は減ります。
60,000画像なら、batch 32は1,875 stepですが、8192は8 stepです。同じ「1 epoch」でも最適化の軌跡は同じではありません。今回の目的は計算性能とmemoryの測定なので、最終精度の比較とは分離します。
実験モデルと測定範囲
モデルは前の記事と同じ105,866パラメータのCNNです。
input 1×28×28
→ Conv 1→16 → ReLU → MaxPool
→ Conv 16→32 → ReLU → MaxPool
→ Linear 1568→64 → ReLU
→ Linear 64→10
条件は次の通りです。
- FashionMNIST学習データ60,000枚
- FP32、TF32は
ieee指定で無効 - Adam、learning rate 0.001
num_workers=0、pin_memory=True- batch：32、64、128、256、512、1024、2048、4096、8192
- 各batchを3回、順序は昇順→降順→昇順
- 毎回同じ初期weight、shuffle seed 42
- 同じbatch shapeで合成1 stepをwarm-upし、測定外
- DataLoader、CPU→GPU転送、forward、backward、Adam更新を含む
- CUDAの非同期実行を考慮して前後に
synchronize()
torch.cuda.synchronize()
started = time.perf_counter()
for images, labels in loader:
images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)
optimizer.zero_grad(set_to_none=True)
loss = loss_function(model(images), labels)
loss.backward()
optimizer.step()
torch.cuda.synchronize()
seconds = time.perf_counter() - started
forwardだけのmicrobenchmarkではありません。研究室で実際に学習scriptを回したときの待ち時間へ近づけるため、データ供給とoptimizerも含めました。
throughputは256付近から頭打ちになった
実データ3回の中央値です。
| batch | images/s | epoch秒 | allocated MiB |
|---|---|---|---|
| 32 | 16,062 | 3.735 | 70.7 |
| 64 | 22,425 | 2.676 | 75.7 |
| 128 | 29,936 | 2.004 | 85.9 |
| 256 | 34,738 | 1.727 | 106.2 |
| 512 | 34,896 | 1.719 | 232.7 |
| 1024 | 37,124 | 1.616 | 399.3 |
| 2048 | 37,332 | 1.607 | 728.9 |
| 4096 | 36,768 | 1.632 | 1,386.0 |
| 8192 | 36,983 | 1.622 | 2,704.3 |
step数は順に1,875、938、469、235、118、59、30、15、8でした。step latencyの中央値は1.99 msから202.80 msへ増えています。reservedを含む全raw値は再現用JSONへ保存しました。
32→256では、step時間は約3.7倍になりましたが、1 stepの画像数は8倍です。そのためepoch全体は約2.16倍速くなりました。
ところが256以降、throughputは約3.5～3.7万images/sで頭打ちです。batchを倍にするとstep時間もほぼ倍になり、epoch時間が縮まらなくなりました。
この先は「もっと大きくすれば速い」領域ではなく、同じthroughputのためにVRAMと1 step latencyを増やす領域です。
実用上のsweet spotは最大値とは限らない
今回の最高中央値はbatch 2048の37,332 images/sでした。しかしbatch 256も34,738 images/sで、最高値の93.1%です。
その差は1 epochあたり約0.12秒です。一方、allocatedは106.2 MiBと728.9 MiBで約6.9倍違います。
空いたVRAMは、より大きいモデル、高解像度入力、data augmentation、別process、推論serverなどへ使えます。またbatchを小さくすると、1 epochあたりのweight更新回数が増えます。最高throughputの一点だけを選ぶのではなく、速度・memory・収束を合わせて選ぶべきです。
16GBを超えたのに、なぜ自然OOMにならなかったのか
次に、DataLoaderを外し、同じCNNへ合成画像を1 batchだけ入力しました。これは巨大batchでmemory境界を見るprobeです。
自然条件ではbatch 131,072までOutOfMemoryErrorになりませんでした。PyTorchのpeak値は次のようになりました。
- batch 32,768：allocated 10,613 MiB、reserved 12,186 MiB、約199,072 images/s
- batch 49,152：allocated 15,887 MiB、reserved 18,244 MiB、約56,761 images/s
- batch 98,304：allocated 22,488 MiB、reserved 27,196 MiB、約20,498 images/s
- batch 131,072：allocated 20,842 MiB、reserved 28,684 MiB、約24,173 images/s
torch.cuda.mem_get_info()が報告したtotalは16,302.6 MiBです。それなのにPyTorchの論理的なpeak counterはそれを超えました。同時にbatch 32,768→49,152でthroughputは約71%低下しています。
この観察から、少なくともPyTorchのcounterが物理VRAMへの常駐量と同じではなく、OOMしないことも高速に処理できることを保証しないと分かります。
このPCはWindowsのWDDM経由でWSL2へGPUを公開しています。memory virtualizationや退避が関係した可能性は高いものの、この実験だけでどのpageがどこへ置かれたかまでは証明できません。NVIDIA公式のCUDA on WSL User Guideも、WSLではGPU機能とmemory利用に固有の制約があると説明しています。
実際のCUDA OOMを安全に再現する
自然OOMを求めてsystem memoryまで使い切るのは、安全な実験ではありません。そこでPyTorch caching allocatorへ、可視memoryの25%を上限として設定しました。
torch.cuda.set_per_process_memory_fraction(0.25, 0)
公式のset_per_process_memory_fractionによると、許容量はtotal visible memory × fractionで、processがそれを超えるallocationを試すとallocatorがOOMをraiseします。
今回の許容量は約3.98 GiBでした。
- batch 8,192：成功、peak allocated 2,703 MiB
- batch 16,384：成功、peak allocated 3,803 MiB
- batch 24,576：
torch.OutOfMemoryError - OOM時の追加要求：588 MiB
- OOM時：PyTorch allocated約3.15 GiB、reserved but unallocated約444 MiB
これは「RTX 5070 Tiの物理限界が24,576」という意味ではありません。教材として再現可能なallocator上限で起こしたcontrolled OOMです。自然条件の挙動と分けて記録しています。
OOM messageにも検算が必要
今回のexception messageには、正しい上限3.98 GiBや追加要求588 MiBとともに、process使用量が非現実的な17179869184.00 GiBという値も含まれました。
16GB GPUでこの値はあり得ません。WSL2、PyTorch 2.13 nightly相当、process fractionを組み合わせたdiagnostic表示上の問題と考えられますが、原因は未確定です。
例外messageを丸ごと信じず、次を合わせて保存しました。
torch.cuda.mem_get_info()のfree / totalmax_memory_allocated()max_memory_reserved()- 設定したprocess fraction
- batch size
- PyTorch / CUDA / driver version
研究では「errorがそう言った」だけでなく、単位と桁がhardware仕様に合うか検算します。
OOMからprocessを回復させる
OOMをcatchしても、失敗した計算graphやTensorへの参照が残っているとmemoryは解放されません。
try:
loss = torch.nn.functional.cross_entropy(model(images), labels)
loss.backward()
optimizer.step()
except torch.OutOfMemoryError as exc:
print(exc)
finally:
del model, optimizer, images, labels, loss
gc.collect()
torch.cuda.empty_cache()
empty_cache()だけでは、参照中のactive Tensorを解放しません。まずPython側の参照を外す必要があります。今回のscriptはOOMをcatchした後もJSONと図を正常に書き出し、processを終了できました。
実際の学習なら、OOM後に次を試します。
- batch sizeを半分にする
- FP16/BF16 AMPを使う
- gradient accumulationで小batchを複数stepまとめる
- 入力解像度・sequence lengthを下げる
- activation checkpointingを使う
- 不要なTensorやloss historyがGPU上に残っていないか調べる
gradient accumulationはOOM対策だが同じ計算ではない
たとえばmicro batch 64を4回累積すれば、optimizerから見たeffective batchは256です。
accumulation_steps = 4
optimizer.zero_grad(set_to_none=True)
for step, (images, labels) in enumerate(loader, start=1):
loss = loss_function(model(images.cuda()), labels.cuda())
(loss / accumulation_steps).backward()
if step % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
ただしBatchNormのstatistics、dropout mask、最後の端数batchなどにより、巨大batchを一度に処理する場合と完全には同じにならないことがあります。今回のCNNにはBatchNormを入れていませんが、一般化するときには注意が必要です。
再現方法
cd /mnt/c/Users/user/Documents/Codex/2026-08-24/ko/work/ai_lab
/opt/ai-lab/venv/bin/python experiment_fashion_batch_oom.py \
--repeats 3 \
--data-dir /opt/ai-lab/data \
--output results/fashion_batch_oom.json \
--figure ../../outputs/fashion-batch-throughput-oom.png
JSONには27回の実データ学習、自然条件の巨大batch probe、controlled OOMの全測定値とexception messageを保存しています。
batch sizeを選ぶためのチェックリスト
- throughputはforwardだけでなく、実際のtraining loopで測る
- CUDA計測の前後で
synchronize()する - warm-upを除き、複数回の中央値と範囲を出す
- batchを増やしたときのoptimizer step数も確認する
- 最高throughputだけでなく、その90～95%を小memoryで得る点を探す
allocated、reserved、物理resident memoryを混同しない- OOMしなくても急な速度低下がないか確認する
- OOM messageの単位・桁をhardware仕様と照合する
- controlled limitと物理限界を明確に区別する
- 精度・収束比較は複数seedと適切なlearning rateで別実験にする
まとめ
FashionMNISTの小型CNNでは、batch 32→256でthroughputが約2.16倍になりましたが、256以降はほぼ頭打ちでした。batch 256は最高値の約93%を106 MiBで達成し、8192は約6.5%速いだけで2.7GBを使いました。
さらにWSL2では、PyTorchのmemory counterが16GBを超えても自然OOMにならず、その代わりthroughputが大きく崩れる領域を観測しました。したがって最適batchは「OOMしない最大値」ではありません。
まず小さいbatchから倍々に増やし、throughputの伸びが止まる点、VRAM、step latency、収束を合わせて選ぶのが実践的です。さらにDataLoaderのnum_workersとpin_memoryを変えると、GPUへデータを供給する側のbottleneckを分離できます。
