Qhapaq Lab

ONNX / TensorRT

ONNX / TensorRT は、深層学習系の将棋AIで学習済みニューラルネットワークを高速に推論するための実装・運用技術である。

dlshogi のような AlphaZero 型エンジンでは、モンテカルロ木探索(MCTS)の葉局面で方策・価値ネットワークを何度も評価する。そのため、学習時の PyTorch モデルをそのまま使うだけでなく、ONNX へ変換し、ONNX Runtime や TensorRT で実行することが重要になる。

ONNXとは

ONNX は Open Neural Network Exchange の略で、ニューラルネットワークをフレームワーク間で交換しやすくするためのモデル形式である。PyTorch で学習したモデルを ONNX に変換しておくと、C++ エンジン、ONNX Runtime、TensorRT などから利用しやすくなる。

コンピュータ将棋の文脈では、ONNX は次のような橋渡しをする。

  • Python / PyTorch で学習した policy/value network を保存する。
  • C++ の USI エンジンから読み込める形にする。
  • ONNX Runtime や TensorRT など、複数の推論バックエンドを選べるようにする。
  • batch size を動的または固定で扱えるようにする。

TensorRTとは

TensorRT は NVIDIA GPU 向けの推論最適化 SDK である。ONNX などから読み込んだネットワークを解析し、GPU 上で高速に実行できる engine を構築する。

TensorRT では、FP32 だけでなく FP16 や INT8 などの精度を使った高速化、レイヤー融合、カーネル選択、optimization profile による dynamic shape 対応などが重要になる。dlshogi のように探索中に大量の局面を評価するエンジンでは、推論速度の改善がそのまま探索速度や自己対局生成速度に影響する。

dlshogiでのモデル変換

DeepLearningShogi には dlshogi/convert_model_to_onnx.py があり、学習済みモデルを ONNX へ変換する。

このスクリプトでは、policy_value_network を作成し、serializers.load_npz で重みを読み込み、torch.onnx.export で ONNX を出力する。入力名は input1, input2、出力名は output_policy, output_value である。

--fixed_batchsize を指定しない場合は、次の 0 次元が batch_size として dynamic axes になる。

  • input1
  • input2
  • output_policy
  • output_value

一方、--fixed_batchsize を指定した場合は固定 batch size の ONNX として出力される。固定 batch は最適化しやすい場合があるが、実際の探索で異なる batch size を使う場合には不便になる。dynamic batch は柔軟だが、推論バックエンド側で optimization profile などの設定が重要になる。

ONNX Runtime版

DeepLearningShogi には usi_onnxruntime があり、ONNX Runtime を使って推論する構成がある。

usi_onnxruntime/nn_onnxruntime.cpp では、Ort::Session で ONNX モデルを読み込み、forward で batch_size を先頭次元に持つ入力テンソルを作る。入力 shape は概念的に次のようになる。

input1: [batch_size, features1, 9, 9]
input2: [batch_size, features2, 9, 9]
output_policy: [batch_size, move_labels]
output_value: [batch_size, 1]

また、Python 側の dlshogi/test.py や dlshogi/utils/hcpe_re_eval.py では、CUDAExecutionProvider や TensorrtExecutionProvider を指定して ONNX Runtime から推論している。ONNX Runtime では、CUDA Execution Provider は NVIDIA GPU 上での汎用実行、TensorRT Execution Provider は TensorRT を使った高速実行を担う。

TensorRT版

USI エンジン側の usi/nn_tensorrt.cpp では、ONNX ファイルを TensorRT parser で読み込み、TensorRT engine を構築する。

主な流れは次の通りである。

1. ONNX parser で network を読み込む
2. builder / config を作る
3. FP16 または INT8 などの精度を選ぶ
4. input1, input2 の optimization profile を設定する
5. serialized engine を構築する
6. engine をファイルへ保存し、次回以降は読み込む
7. forward 時に batch_size を設定して enqueue する

DeepLearningShogi の TensorRT 実装では、onnx_filename + gpu_id + max_batch_size + .serialized のようなファイル名で serialized engine を保存する。既に serialized engine があれば読み込み、なければ ONNX から build して保存する。これは、TensorRT engine の build が重い処理であるためである。

また、onnx_filename + .calibcache が存在する場合は INT8 を使い、そうでなければ FP16 が使える環境では FP16 を使う、という分岐がある。INT8 は速くなりうる一方、calibration data と calibration cache の整合が重要になる。

バッチ推論との関係

ONNX / TensorRT はバッチ推論と強く関係する。

探索部は複数の未評価局面を集め、DNN_Batch_Size や policy_value_batch_maxsize に応じた batch として推論バックエンドへ渡す。ONNX では batch 次元を dynamic axes として持てる。TensorRT では optimization profile の MIN, OPT, MAX shape によって、実行時に許される batch size の範囲を定める。

そのため、次の値は互いに整合していなければならない。

  • ONNX の batch 次元が固定か dynamic か。
  • TensorRT engine の max batch size。
  • USI オプションの DNN_Batch_Size。
  • 自己対局生成時に指定する batchsize。
  • GPU メモリ、モデルサイズ、推論精度。

運用上の注意

ONNX / TensorRT 周辺は、将棋エンジン本体とは別の意味で環境依存が強い。

  • CUDA、cuDNN、TensorRT、ONNX Runtime のバージョン整合が必要である。
  • TensorRT engine は GPU、TensorRT バージョン、precision、profile、モデルに依存し、安易に使い回せない。
  • モデルを差し替えたら serialized engine や calibration cache を作り直す必要がある。
  • dynamic batch を使う場合は、実際に使う batch size が optimization profile の範囲に入っている必要がある。
  • INT8 は calibration data に依存するため、速さだけでなく棋力や評価のずれも確認する必要がある。
  • ONNX Runtime で TensorRT Execution Provider を使う場合、TensorRT が対応できない subgraph は CUDA や CPU へ fallback することがある。

関連項目

参考にしたホームページ