Pytorch DataLoaderでまだまだ高速化 #
初めに #
皆さんはDataLoaderを正しく使えていますか？
「データの読み込みで時間がかかる」，「GPUが動いてない時間がある」などの問題を解決できるかもしれません．
今回はこれらの問題を解決するDataLoaderの使い方を解説いたします．
DataLoaderの使い方 #
loader = DataLoader(ds, batch_size=32, shuffle=True)
基本的な使い方としてDataLoaderをこのように使用していると思います．
しかしこれだとDataLoaderの能力を完全には発揮できていません．
DataLoaderの引数 #
今回使用するDataLoaderの引数を紹介します．
| 引数 | 型 | 内容 |
|---|---|---|
| dataset | Dataset | データセット |
| batch_size | int | 1バッチあたりのサンプル数 |
| shuffle | bool | エポック毎にシャッフル |
| num_workers | int | 並列データ読み込みプロセス数 |
| pin_memory | bool | CPU→GPU転送を非同期に（CUDA環境ならTrue） |
| drop_last | bool | 最後に余る端数バッチを捨てる |
| persistent_workers | bool | エポック間でプロセスを維持 |
| prefetch_factor | int | バッチの先読み数 |
今回の重要人物 #
num_workers #
バッチデータを読み込むプロセス数を指定できます．データを並列に読み込むことで高速化することができます．
ただし，プロセスを起動するので，CPUと相談してください．
persistent_workers #
通常，DataLoaderが最後のバッチまでの処理を行うとプロセスを終了します．そうすると次のエポック開始時にプロセスを再生成する必要があります．そこで，persistent_workersをTrueとすることで，プロセスを終了させず，次のエポックに再利用することができます．
prefetch_factor #
各プロセスが先読みするバッチ数を指定できます．通常は1つのプロセスにつき2つのバッチしか読み込んでおらず，バッチデータを投げ込んで処理している間待機状態となります．この時間で次のバッチを読み込み，空いている時間を効率よく使うことができます．
ただし，バッチデータを多く保持するので，メモリのオバーフローに注意してください．
使用例 #
今回の紹介した引数を使用したDataLoaderの使用例です．
DataLoader(ds,
batch_size=32,
shuffle=True,
num_workers=4,
pin_memory=True,
prefetch_factor=4,
persistent_workers=True,
drop_last=True)
この場合，4つのプロセス（num_workers=4）が4バッチ先読み（prefetch_factor=4）しているので16バッチがキューに入ることになります．
サンプル数に換算すると，4×4×32で512サンプルをメモリ上にスタンバイすることになります．
このため，メモリの占有率が上昇しますが，読み込みに時間がかかるデータにおいては，先に読み込むことでプロセスが待機状態になることを防ぐことができます．
まとめ #
今回は，PytorchのDataLoaderで処理を高速化する方法について紹介しました．
皆さんは正しく使えていましたか？
使えていなかった方はこれを機に学習の高速化を目指しましょう．
