TadaoYamaokaの開発日記

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

dlshogiで量子化認識トレーニングを試す

先日、PyTorchでFXグラフモードで量子化認識トレーニングを試して、精度が大きく落ちることなく、CPUでの推論が速くなることを確認した。

そこで、dlshogiのモデルで量子化認識トレーニングを試して、CPUで強くなるかを検証した。

量子化認識トレーニン

先日の記事にも書いたが、量子化認識トレーニングは、訓練時に量子化(浮動小数点を整数にする)のスケールパラメータを学習することで、精度を保ったまま整数化により推論を高速する手法である。

PyTorchのFXグラフモードを使うことで、比較的簡単に量子化認識トレーニングを実装できる。

ONNX化の課題

dlshogiのモデルを量子化認識トレーニングすると、ONNXに出力する際、

Exporting the operator 'quantized::batch_norm2d_relu' to ONNX opset version 14 is not supported.

というエラーが発生し、ONNXに変換できなかった。

量子化した後のTorchScriptモデルで使われているquantized::batch_norm2d_reluオペレータが、ONNXに対応していないことが原因のようだ。

ONNX化はあきらめて、LibTorchで推論することにした。
LibTorchでも、TorchScriptモデルを使えば通常モデルに比べて高速に推論することができる。

強さ測定

CPUで実行することを考慮して、dlshogiの30ブロック384フィルタのモデルから5ブロック96フィルタおよび20ブロック256フィルタのモデルに蒸留を行った。

量子化前のモデルも含めて、開始局面から16手ランダムあり、1手2秒で連続対局を行った。
ハードウェアは、CPU 4コア(Xeon E3-1275 v6)、メモリ 64GBのマシンを使用した。
バッチサイズは2、探索スレッド数は1、ルートでの詰み探索あり、PVの詰み探索ありとした。

   # PLAYER             :  RATING  ERROR  POINTS  PLAYED   (%)  CFS(%)    W    D    L  D(%)
   1 dlshogi_20bqat     :   280.7   49.2   472.0     800    59     100  465   14  321     2
   2 dlshogi_5bfp32     :   122.9   56.8   273.5     802    34     100  268   11  523     1
   3 dlshogi_5bqat      :    34.8   65.0   106.5     402    26      85  106    1  295     0
   4 dlshogi_20bfp32    :     0.0   ----    59.5     400    15     ---   57    5  338     1

White advantage = 49.19 +/- 9.52
Draw rate (equal opponents) = 1.45 % +/- 0.31

※同系列で偏らないように基準ソフトとして3コマ関係のソフトも加えているが省略している。
※末尾fp32が量子化前のモデル、qatが量子化したモデル

5ブロックのモデルは、量子化するとかえって弱くなった。
20ブロックのモデルは、量子化により強くなっている。
量子化前は、5ブロックの方が強い。

精度

各モデルのfloodgateの棋譜に対する正解率は以下の通り。

モデル 方策正解率 価値正解率
5ブロック(FP32) 0.4570605 0.7262039
5ブロック(量子化) 0.4400687 0.7097659
20ブロック(FP32) 0.5414595 0.7704120
20ブロック(量子化) 0.5317605 0.7590858

5ブロックのモデルでは、量子化により方策、価値ともに正解率がそれぞれ1.7%と1.6%低下している。
20ブロックのモデルでは、量子化により方策、価値ともに正解率がそれぞれ1.0%と1.1%低下している。
20ブロックの方が量子化による正解率の低下が小さい。

NPS

平手開始局面で、10秒探索させた際のNPSは以下の通り。

モデル NPS
5ブロック(FP32) 2049
5ブロック(量子化) 2453
20ブロック(FP32) 38
20ブロック(量子化) 165

5ブロックのモデルでは、量子化によりNPSが1.2倍になっている。
20ブロックのモデルでは、量子化によりNPSが4.34倍になっている。

20ブロックのモデルは、量子化による速度向上の効果が高い。
これは、モデルの推論時間が探索の多くを占めるためと思われる。

20ブロックのF32モデルは、NPSが38しかでておらず、少ない探索ではMCTSの結果が安定しないため、比較したモデルで一番弱かったのも納得できる。

floodgateに放流

比較したモデルで最もレーティングが高かった20ブロックの量子化モデルをfloodgateに放流して、レーティングを計測した。

floodgateのレーティングは約3600であった。

NNUE系は2コアでも約3800であるので、CPUではまだNNUE系を上回っていない。

まとめ

dlshogiで量子化認識トレーニングを試して、量子化前後の強さを比較した。
結果、20ブロックのモデルで量子化の効果が大きいことが確認できた。
5ブロックでは推論は速くなるものの精度低下により逆に弱くなった。

dlshogiを30ブロックの最新モデルを20ブロックに蒸留して量子化すると、4コアのCPUでNPSが165しかないがレーティングが約3600あることがわかった。
モデル精度を向上することで、CPUでもNNUE系を上回ることができるか引き続き検証したい。

2023/6/15追記

対戦数232でレーティングは3719となった。