dlshogiのモデルの訓練に使用しているPythonスクリプトは、ベタなforループで記述しており、ミニバッチ作成部分と、ニューラルネットワークの学習の処理をシーケンシャルに実行しており並列化は行っていなかった。
ミニバッチデータの作成は、盤面から入力特徴量を作成する処理が比較的重いため、C++で実装して高速化している。
それでも、ある程度CPU時間を消費している。
一方ニューラルネットワークの学習処理は、GPUで処理しているため、CPUは遊んでいる状態になる。
その間に、次のミニバッチデータの作成を行えば、よりCPUの使用効率を上げることができる。
PyTorchには、データローダーを並列化する仕組みがあるが、マルチプロセスで実装されており、Windowsの場合は、プロセスにデータをpickleで送信する必要があり、データ量が大きいと効率が上がらない。
torch.utils.data — PyTorch 1.10 documentation
なお、Linuxの場合は、マルチプロセスはfork()で実装されているため、プロセス間でメモリが共有されるため転送がいらない。
そこで、Windowsでも高速化できるように、マルチスレッドのデータローダーを実装した。
実装方法
学習処理の間にミニバッチの作成を別スレッドで行うようにする。
実装は単純で、以下のように処理する。
- ミニバッチを作成した後に、次に作成するミニバッチを別スレッドで作成開始するようにする(pre fetch)。
- 次にミニバッチが必要になったタイミングでは、別スレッドで作成していたデータが作成済みになっているのでそれを取得する。
- 以上を繰り返す。
別スレッドでの実行と、完了を待つ処理は、PythonのThreadPoolExecutorを利用して、Futureパターンで実装した。
別スレッドで実行中に、GILを解放しないと、Pythonではマルチスレッドの効果がでないため、C++側でGILを解放する処理を入れた。
また、今までミニバッチ作成の度、メモリをnp.emptyで確保していたが、データローダ作成時に一回作成してそれを使いまわすようにした。
その際、pin_memoryも有効にした。
ソース
データ読み込みの高速化 · TadaoYamaoka/DeepLearningShogi@5d50e8b · GitHub
DataLoaderのマルチスレッド化 · TadaoYamaoka/DeepLearningShogi@865a199 · GitHub
※git pullした場合はcppshogiのリビルドが必要なため注意
測定結果
変更前後で、学習時間は以下の通りになった。
- 1,439,175局面の学習時間を測定
- SWAのbn_updateとテストの時間含む
| 時間(時:分:秒) | 比 | |
|---|---|---|
| 変更前 | 0:06:33 | 100% |
| 変更後 | 0:04:47 | 73.0% |
学習時間を73%に短縮することができた。
なお、np.emptyで毎回メモリ確保するのをやめるだけでも、学習時間は84%になった。
メモリ確保も比較的重かったようだ。
まとめ
Windowsでも高速で動くマルチスレッドのデータローダを実装した。
その結果、学習時間を73%に短縮することができた。
