TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

将棋AIの進捗 その58(訓練の省メモリ化と高速化)

現状のdlshogiのモデルの訓練の実装では、訓練データが多い場合にメモリを大量に消費する課題があった。

dlshogiのモデルの訓練は、訓練データをすべてメモリに読み込んで処理を行っている。
そのため、一度に学習できるデータはメモリに読み込める分に制限される。
それ以上のデータを学習する場合は、訓練データを分割して学習する必要がある。

また、1台のマシンでGPUを複数使用して同時に実験を行いたい場合、それぞれメモリを消費するため、メモリが不足する。

対策

メモリ消費を抑えるために、前処理を行ったデータをキャッシュファイルに出力して、キャッシュファイルをシークしながらランダムアクセスして読み込む機能を追加した。
OSがページファイル単位でメモリに読み込んでキャッシュするため、物理メモリに乗る容量であれば、ほとんど処理時間は変わらない。
また、複数の実験を同時に行う場合には、同一のキャッシュファイルを使用すると、OSが同じページファイルから読み込むのでメモリ消費を抑えられる。

高速化

キャッシュファイルから読み込む場合、IO待ちが発生するため、その間にCPUを遊ばせないように、ミニバッチを作成する際のデータ読み込みから特徴量を作成する一連の繰り返し処理をOpenMPで並列で処理するようにした。
キャッシュファイルを使用しない場合でも、特徴量作成処理がわずかにボトルネックになっていたため、同様にOpenMPで並列化を行った。
スレッド数2以上にしてもかえって遅くなったため、スレッド数は2とした。

測定

キャッシュ機能と並列化を実装して、学習時間を比較したところ、10ブロックモデル・4096バッチ・1000ステップ学習の平均学習時間(GPUはA100)は、

条件 平均学習時間(h:m:s)
キャッシュファイルなし(並列化なし) 0:03:36
キャッシュファイルなし(並列化あり) 0:03:27
キャッシュファイルから 0:03:32

となった。

並列化により、元のバージョンから4.01%速くなった。
キャッシュファイルから読み込む場合は、元のバージョンから1.85%速くなった。並列化ありでキャッシュファイルを使用しない方が速いが、元のバージョンより遅くはなっていないので良しとする。

使い方

train.pyのオプションで、--cacheでキャッシュファイルを指定すると、キャッシュファイルがない場合は作成され、キャッシュファイルがある場合はキャッシュファイルを使用する。
キャッシュファイルがない場合は、いったんデータをすべてメモリに読み込んでからキャッシュファイルを作成した後、メモリを解放してから、キャッシュファイルを使用する。

まとめ

1台のサーバで、複数実験を行う際にメモリ不足が起きていたため、キャッシュファイルから読み込むことで省メモリ化するオプションを実装した。
また、IO待ちが発生する対処として特徴量作成処理の並列化も実装したことで、キャッシュファイルを使用しない場合も4%ほど高速化された。