PyTorch DataLoaderは何workerが速い？WSL2でnum_workers・pin_memoryを実測

実験日: 2026-08-24 / PyTorch 2.13.0+cu130
CPU: Core i7-14700F（28 logical CPU）/ GPU: RTX 5070 Ti / WSL2

GPU学習が遅いとき、原因がGPUとは限りません。CPUが画像を読み、transformし、batchにまとめる間、GPUが次のデータを待っていることがあります。

GPU使用率だけを見て判断すると、データ供給側の遅さを見落とします。ここが少し厄介です。

FashionMNIST CNNをnum_workers=0, 1, 2, 4, 8, 12, 16で実際に学習し、軽いToTensorだけの場合と、CPU augmentationを加えた場合を比較しました。worker起動を含む最初のepochと、persistent_workers=Trueでprocessを維持した2 epoch目も分けています。

結果は、augmentationありで0 workerの約8,830 images/sに対し、12～16 workerの反復中央値は約62,100～63,800 images/sでした。約7倍です。一方、pin_memoryは各5回の中央値でOFF 1.634秒、ON 1.635秒となり、今回のend-to-end学習では差を確認できませんでした。

worker数の結論

軽いtransform（ToTensorのみ）

  • 0 worker：30,268 images/s
  • 2 worker：60,881 images/s
  • 4 worker：108,782 images/s
  • 8 worker：116,828 images/s
  • 12 worker：119,045 images/s
  • 16 worker：117,262 images/s
  • このsweepでは12付近でplateau

CPU augmentationあり

  • 0 worker：8,830 images/s
  • 2 worker：20,626 images/s
  • 4 worker：41,292 images/s
  • 8 / 12 / 16 workerは追加で各5回測定
  • 反復中央値：8は55,196、12は62,103、16は63,753 images/s
  • 12→16の差は約2.7%。12～16をplateau領域と判断

pin_memory、augmentation、4 worker、各5回

  • OFF：中央値1.634秒（1.611～1.649秒）
  • ON：中央値1.635秒（1.599～1.669秒）
  • 範囲が大きく重なり、end-to-end差なし
WSL2 PyTorch DataLoaderでtransform負荷別のnum_workers、worker起動費、persistent worker、pin_memoryを実測
WSL2 PyTorch DataLoaderでtransform負荷別のnum_workers、worker起動費、persistent worker、pin_memoryを実測

図の左側は温間epochです。augmentationの8/12/16だけは各5回の中央値、その他は全条件sweepの値です。右側は同じsweepで、最初のepochとpersistent workerを使った2回目を比較しています。

DataLoaderは何を並列化するのか

Dataset.__getitem__()は、indexを受け取って1 sampleを返します。画像fileの読み込み、decode、crop、rotation、Tensor化などは通常ここで実行されます。

num_workers=0では、main processが学習loopとdata preparationの両方を順番に行います。

main: [load/transform] → [GPU training] → [load/transform] → [GPU training]

num_workers>0では、subprocessが次のbatchを準備し、main processがGPU計算を進めます。

worker 0: [prepare batch A] ───── [prepare batch C]
worker 1: [prepare batch B] ───── [prepare batch D]
main/GPU:        [train A][train B][train C][train D]

理想的にはdata preparationとGPU computeが重なります。ただしprocess作成、inter-process communication（IPC）、queue、memory copyにもcostがあります。workerを増やせば常に速いわけではありません。

PyTorch公式のPerformance Tuning Guideも、num_workers>0でasynchronous loadingとaugmentationを重ねられる一方、workload、CPU、GPU、data保存場所に応じて調整すべきだと説明しています。

2種類のworkloadを用意した理由

worker数の効果はDatasetが行う仕事で変わります。

軽い条件です。

transform = transforms.ToTensor()

augmentation条件です。

transform = transforms.Compose([
    transforms.RandomResizedCrop(28, scale=(0.75, 1.0)),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
])

FashionMNISTは小さい28×28画像で、fileもlocal SSD上にあります。それでもrotationとcropを60,000枚へ毎epoch適用するとCPU側の比率が増えます。

「DataLoaderは4 workerが最適」と一般化するのではなく、自分のdecode・transform・storageを含むpipelineを測る必要があります。

実験条件

  • FashionMNIST学習画像60,000枚
  • CNN 105,866パラメータ
  • batch 256、Adam、FP32、TF32無効
  • num_workers：0 / 1 / 2 / 4 / 8 / 12 / 16
  • prefetch_factor=2（workerありの場合）
  • persistent_workers=True（workerありの場合）
  • pin_memory=Trueを基本条件
  • 同一loaderで2 epoch
  • epoch 1：worker process起動とinitial prefetchを含む
  • epoch 2：workerを維持した温間
  • DataLoader、host→device転送、forward、backward、Adamをすべて計測
  • GPU/libraryは事前に合成1 stepでwarm-up
  • CUDA計測の前後でsynchronize()
loader = DataLoader(
    dataset,
    batch_size=256,
    shuffle=True,
    num_workers=workers,
    pin_memory=True,
    persistent_workers=workers > 0,
    prefetch_factor=2 if workers > 0 else None,
)

persistent_workersとprefetch_factorはworker=0では使えません。実装時は条件分岐が必要です。

軽いtransformでも4 workerまで大きく伸びた

ToTensorだけの温間throughputは、0 workerの30,268から4 workerの108,782 images/sへ約3.59倍になりました。

FashionMNISTは小さくても、Pythonでsampleを1件ずつ取り出し、PIL imageをTensorへ変換し、batchへcollateする処理があります。小型CNNのGPU計算が速いため、CPU側のわずかな仕事が相対的にbottleneckになりました。

8 workerは116,828、12は119,045、16は117,262 images/sです。4→12の増加は約9.4%ですが、12→16では少し低下しました。

したがって軽い条件では4がcost performanceのよい候補、最高throughputを狙うなら8～12付近です。16 processを常時維持する利益はありませんでした。

augmentationありでは並列化の効果がさらに大きい

augmentation条件のsweepです。

workers epoch 1秒 warm秒 warm images/s
0 6.327 6.795 8,830
1 5.792 5.757 10,422
2 2.967 2.909 20,626
4 1.493 1.453 41,292
8 0.903 0.852 70,398
12 0.830 0.853 70,322
16 1.055 0.937 64,050

単発sweepでは8と12がほぼ同じで、16は低下しました。しかしepochが1秒未満になるとOS schedulerやprocess状態の揺らぎが順位へ影響します。別のsweepでは12が最高になりました。

そこで8、12、16だけを順序ローテーションして各5回追加測定しました。

  • 8 worker：55,196 images/s（54,306～55,660）
  • 12 worker：62,103 images/s（61,716～62,660）
  • 16 worker：63,753 images/s（62,614～63,974）

絶対値はsweepと追加反復で変わりましたが、0～4から大きく改善し、12～16で伸びが小さくなる傾向は共通です。16は12より約2.7%速いだけなので、このPCでは12付近を実用的なplateau開始点と判断しました。

なぜworker=1の最初だけ遅いのか

軽い条件のworker=1は、epoch 1が3.371秒、温間epochが1.910秒でした。0 workerの約1.97秒より初回だけ大幅に遅くなっています。

1 workerではdata preparationをmain processから分離できますが、並列度は増えず、最初に次のcostを払います。

  • subprocess作成
  • Python moduleとDataset状態の準備
  • queueとshared-memory経路の準備
  • 最初のbatchのprefetch

persistent_workers=Trueにすると、Datasetを1周した後もworkerを終了しません。2 epoch目以降はprocess起動を繰り返さずに済みます。

ただしworkerを維持するとprocessとDatasetのmemoryも維持されます。epoch間に長時間止まるjobや、巨大なDataset objectではtrade-offになります。

prefetch_factor=2は何を意味するか

workerがあるとき、各workerは先のbatchをqueueへ準備します。prefetch_factor=2なら概念上、workerごとに最大2 batchを先読みします。

16 workerなら最大32 batch分がpipeline上に存在し得ます。batch 256の高解像度画像ならhost memoryやshared memoryを多く使います。

workerを増やしてBus error、shared memory不足、RAM圧迫が起きる場合は、worker数、batch size、prefetch_factorを下げます。throughputだけでなく/dev/shmとRAMも観測対象です。

pin_memoryは速くならなかった

pin_memory=TrueはDataLoaderがpage-locked host memoryへbatchを置く設定です。pinned memoryからGPUへのcopyは、non_blocking=Trueと組み合わせてasynchronousにできます。

images = images.cuda(non_blocking=True)
labels = labels.cuda(non_blocking=True)

PyTorch公式のpin memory / non_blocking guideも、DataLoader側でpinningし、non-blocking transferを使う方法を解説しています。

しかし今回のaugmentation・4 worker・各5回では次の結果でした。

  • OFF：1.634秒（1.611～1.649）
  • ON：1.635秒（1.599～1.669）

中央値はほぼ同じで範囲も重なります。28×28×1 channel、batch 256は転送量が小さく、CPU transformと学習計算の方が支配的です。また専用CUDA streamで明示的にprefetchしていないため、転送と計算のoverlap余地も限定的です。

これはpin_memoryが無意味という結論ではありません。高解像度・大batchでH2D転送が大きい場合は効果が変わります。pinned memoryはswapできないため、過剰使用はhost RAMへ圧力をかける点にも注意します。

worker数を決める実践手順

  1. 実際のDatasetとtransformを固定する
  2. 0、1、2、4、8のように増やす
  3. 最初のepochと温間epochを分ける
  4. throughputが大きく伸びなくなる点を探す
  5. その周辺だけ複数回測る
  6. CPU使用率、RAM、/dev/shmも確認する
  7. augmentationやstorageを変えたら再測定する

「論理CPU数と同じworker」にする必要はありません。今回も28 logical CPUに対し、軽負荷は8～12、augmentationは12付近からplateauでした。

再現方法

cd /mnt/c/Users/user/Documents/Codex/2026-08-24/ko/work/ai_lab

/opt/ai-lab/venv/bin/python experiment_dataloader_bottleneck.py \
  --data-dir /opt/ai-lab/data \
  --output results/dataloader_bottleneck.json \
  --figure ../../outputs/dataloader-bottleneck-workers16.png

図だけ再作成する場合です。

/opt/ai-lab/venv/bin/python plot_dataloader_bottleneck.py \
  results/dataloader_bottleneck.json \
  ../../outputs/dataloader-bottleneck-workers16.png

JSONには全worker sweep、pin_memory各5回、高worker各5回のraw値を保存しています。

まとめ

このPCのFashionMNIST CNNでは、DataLoaderの0 workerが明確なbottleneckでした。軽いtransformでも4 workerで約3.59倍、CPU augmentationありでは12～16 workerが0 workerの約7倍に達しました。

一方、workerを増やす効果はplateauになり、process起動、IPC、prefetch memoryのcostもあります。pin_memoryも今回の小画像ではend-to-end差を確認できませんでした。

結論は、設定値を暗記するのではなく、実際のDataset×transform×storage×modelで、初回と温間を分けて測ることです。画像特徴量HOG+SVMとCNNを同じFashionMNISTで比較すれば、古典機械学習とdeep learningの違いも同じ条件で確認できます。

関連記事