(最新SSD IOはPCIe x4でした。ご指摘ありがとうございます。）

はじめに

Kaggle Advent Calendar 2022 8日目です。

突然ですが、あなたはDNN学習時にどの処理で学習速度が律速しているか把握してますか？

DNN学習には図に示すように大きく3つの要素があります：

（SSDからの）データ読み込み
（CPUによる）データ前処理
（GPUによる）DNN計算

学習時のデータの流れとしては

SSDからデータが読み込まれ、CPUに送られる(SATA or PCIe)
CPUにてaugmentationや正規化などの前処理が行われ、GPUにデータが送られる(PCIe x16)
GPUにてDNNの計算・学習・パラメータ更新が行われる (GDDR)

というのが一般的かと思います。

ここで各データ通信に使われるインターフェース（IO)に注目してみましょう。

まずGPUとGPUメモリ（GPU内DRAM)間通信はGDDRというメモリIOが使われ、その速度は非常に早いです（A6000なら600GB/s)。またCPU-DRAM間もDDR4といったメモリIOで数10GB/s出ており、CPU-GPU間の通信はPCIeが16束になったPCIe x16というIOでこれも高速です。対してSSD-CPU間は普通のPCならSATA、最新PCならばPCIe x4で接続されるのが一般的です。

特にSATAは他のDDR系に対して最悪2桁ほど通信が遅く、注意しないとデータ読み込みで容易に速度が律速してしまうのがわかるかと思います。またHDDを使っているとSSDより読み出し速度は1/3程度になってしまうので使うのは問題外です・・

GPU処理に関する高速化などの記事は大量にある反面、データ読み込みやデータ前処理の高速化の速度律速についてはあまり語られていない気がします。

この記事を読むことで

自分の学習が何で律速しているか把握できるように
データ読み込みや前処理の高速化
学習高速化してより早く多く実験を回す

ようになるきっかけとなれば良いと思います。

どの処理で律速しているか調べる

僕は学習を回したら

1) GPU使用率を調べる
2) CPU使用率を調べる

の2つを見ます。

まずステップ1といてnvidia-smiを叩いてみましょう。

~~このGPU使用率が具体的に何を指しているかは今回省きますが、~~ ここでGPU使用率が100%になっていればGPU律速になっており、データ読み込みやデータ前処理に問題はありません。

ただこの使用率が20~60%と低い値になっている場合、GPUは半分以上の時間データを待っているだけであり、データ読み込みかデータ前処理に処理が律速されていると考えられます。

次にステップ2として、htopなど叩いてみてCPU使用率を見てみましょう（windowsならタスクマネージャーで見れます）

CPU利用率が高い場合（全コア使用されている）、データ前処理で処理が律速されています（CPU律速）。 kaggle notebookはCPUが貧弱（2コアしかない）ため、頻繁にCPU律速が起きてしまいます。ただpytorch dataloaderでnumber_of_workerを設定していない場合は全コアを使用してくれません。自分のコア数程度にworker数を設定して前処理速度を高めましょう。

一方でGPU律速でもCPU律速でもない場合、データ読み込みが全体を律速している可能性があります。上記でも書いた通りSSD⇨CPUのIOは貧弱であり、あまり多くのデータが送れません。そのためCPUやGPUが十分に早いとCPU/GPUのデータ待ちといった状況が発生します（SSD律速）。

各律速の改善方法については次章に書きます。

各処理の速度改善方法

データ読み込み速度の改善

データ読み込みは他IOに比べ数桁遅いため、ボトルネックになりやすいです。

これを改善するためには

読み込むデータのサイズを予め小さくしておく
データを圧縮しておく
DRAMに予めデータを展開してDRAMからデータを読む
早いSSDに買い替える

の4通りが考えられます。まず1つ目ですが、例えばKaggleから配布された画像サイズが1024x1024であった場合でも学習に用いるサイズは512x512と低解像度である場合がほとんどです。その場合、予め画像データを512x512等縮小したものを保存しておき学習時にはそちらを読み込むことで、この例ではデータの読み込みスループットは4倍向上します。

また画像であれば可逆圧縮であるpngに対し、非可逆圧縮であるjpgに保存フォーマットを変えることで大きくファイルサイズを削減することができます。生データであれば（CPU負担は増えますが）npzといった圧縮形式で保存しておくことでデータ読み出しのスループットは向上できます。

またDRAMに自信ニキの場合、予めメインメモリ上に全てのデータを展開しておく力技も考えられます。そうすればDDRでデータ読み込みを行えるため、ボトルネックは解消できます。(書くの大変ですが）

Dataloader resets dataset state - #6 by bfeeny - PyTorch Forums

それでも解消しない場合はPCIe対応のSSDを買いましょう！現在は第四世代PCIeのSSDも売っており、10倍ほど早くなります。マザボが対応しているかはご確認ください。

価格.com - インターフェイス:PCI-Express Gen4のSSD 比較 2022年人気売れ筋ランキング

データ前処理速度の改善

worker数を最大化する
画像を縮小する
CPUを買い替える

のが考えられるアプローチです。

まずデータ前処理が律速している場合、まずはdataloaderのworker数を最大化しましょう。こちらの記事のdataloder編が詳しいです。

qiita.com

また画像の前処理（というか画像処理一般）は画素数をNとした場合O(N²)と書けるため、予め画素数を縮小しておくおくことで前処理速度も大きく向上します。予めresizeした画像を使うのはもちろん、例えばresizeを最初にしてから他前処理をすることでも処理速度の向上が期待できます。

最後にはCPUをコア数多いものに買い替えましょう（8コアあれば大抵十分な気もしますが・・）

GPU処理速度の改善

書いてる人多いので簡単に。

AMPなどFP16・TF32を用いて計算する
pytorch2.0のcompileを使う
batch数を増やす
~~GPUを買い替える~~

aru47.hatenablog.com

AMP(mixed precison）で学習するとFP32に比べメモリ転送量が半分になる他、TensorCoreを使うことで計算速度が大幅に向上します。まず最初に試しましょう。

We just introduced PyTorch 2.0 at the #PyTorchConference, introducing torch.compile!

Available in the nightlies today, stable release Early March 2023.

Read the full post: https://t.co/DFO0HojkMt

🧵below!

1/5 pic.twitter.com/ejEQsJHuB2
— PyTorch (@PyTorch) 2022年12月2日