TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

将棋AIの進捗 その56(データローダーの並列化)

dlshogiのモデルの訓練に使用しているPythonスクリプトは、ベタなforループで記述しており、ミニバッチ作成部分と、ニューラルネットワークの学習の処理をシーケンシャルに実行しており並列化は行っていなかった。

ミニバッチデータの作成は、盤面から入力特徴量を作成する処理が比較的重いため、C++で実装して高速化している。
それでも、ある程度CPU時間を消費している。

一方ニューラルネットワークの学習処理は、GPUで処理しているため、CPUは遊んでいる状態になる。
その間に、次のミニバッチデータの作成を行えば、よりCPUの使用効率を上げることができる。

PyTorchには、データローダーを並列化する仕組みがあるが、マルチプロセスで実装されており、Windowsの場合は、プロセスにデータをpickleで送信する必要があり、データ量が大きいと効率が上がらない。
torch.utils.data — PyTorch 1.7.1 documentation
なお、Linuxの場合は、マルチプロセスはfork()で実装されているため、プロセス間でメモリが共有されるため転送がいらない。

そこで、Windowsでも高速化できるように、マルチスレッドのデータローダーを実装した。

実装方法

学習処理の間にミニバッチの作成を別スレッドで行うようにする。
実装は単純で、以下のように処理する。

  1. ミニバッチを作成した後に、次に作成するミニバッチを別スレッドで作成開始するようにする(pre fetch)。
  2. 次にミニバッチが必要になったタイミングでは、別スレッドで作成していたデータが作成済みになっているのでそれを取得する。
  3. 以上を繰り返す。

別スレッドでの実行と、完了を待つ処理は、PythonのThreadPoolExecutorを利用して、Futureパターンで実装した。
別スレッドで実行中に、GILを解放しないと、Pythonではマルチスレッドの効果がでないため、C++側でGILを解放する処理を入れた。


また、今までミニバッチ作成の度、メモリをnp.emptyで確保していたが、データローダ作成時に一回作成してそれを使いまわすようにした。
その際、pin_memoryも有効にした。

測定結果

変更前後で、学習時間は以下の通りになった。

  • 1,439,175局面の学習時間を測定
  • SWAのbn_updateとテストの時間含む
時間(時:分:秒)
変更前 0:06:33 100%
変更後 0:04:47 73.0%

学習時間を73%に短縮することができた。

なお、np.emptyで毎回メモリ確保するのをやめるだけでも、学習時間は84%になった。
メモリ確保も比較的重かったようだ。

まとめ

Windowsでも高速で動くマルチスレッドのデータローダを実装した。
その結果、学習時間を73%に短縮することができた。