機械学習周りのプログラミング中心。 イベント情報
ポケモンバトルAI本電子書籍通販中

TensorRTのバッチサイズチューニング【コンピュータ将棋】

前回の記事ではNVIDIA GPU上でDNNモデルを高速実行できるライブラリであるTensorRTの使い方について紹介しました。

select766.hatenablog.com

その中で、最適化プロファイルの設定というチューニング項目の効果について紹介します。

TensorRTはモデルとGPUの組み合わせに対して実行計画を最適化しますが、バッチサイズによって適切な実行計画は異なる可能性があります。そのため、複数の実行計画(プロファイル)を使い分ける機能が搭載されています。 プロファイルは、(kMIN=最小バッチサイズ, kOPT=最適バッチサイズ, kMAX=最大バッチサイズ)という3つの数値の組を与えて生成します。最適バッチサイズの入力が与えられたときの実行速度が最大となるように実行計画が最適化されます。

実際にこのパラメータにより実行速度がどう変化するのか実測値を掲載します。実験環境はAWSのp3.2xlargeインスタンスで、Tesla V100 GPUが1基搭載されています。ベンチマークには自作のプログラムを用いました。 https://github.com/select766/tensorrt-shogi

今回使用したモデルは将棋用の9x9の画像を入力とするCNNで、チャンネル数256、161層あります。詳しい定義へのリンクを載せておきます。

https://github.com/select766/neneshogi_v3/blob/b9cf80249baa819673ecba0a0720a3aebaccf7fb/neneshogi/neneshogi/models/resnet.py#L27

ch: 256
depth: 79
block_depth: 2
move_hidden: 2048

最適化プロファイルの前に、32bit浮動小数点数での演算と16bit浮動小数点数での演算での比較を掲載します。最適化プロファイルは、(1, 256, 256)の1つだけを使用します。なお最適化プロファイルにはバッチサイズのほかにメモリ容量という項目もあるのですがこちらは設定していません。

まずは1バッチあたりの計算時間を比較します。

f:id:select766:20200502093159p:plain
バッチサイズと計算時間

理論的な計算量はバッチサイズに比例しますが、実際の計算時間は階段状の変化を見せています。バッチサイズが一定の幅に入っている場合に計算時間はほぼ一定になっているとみなすことができます。また浮動小数点数の精度により大きな速度差があり、バッチサイズ250の時には16bitは32bitの6.5倍速にもなっています。計算誤差が許容できるかどうかには拠りますが、16bitで計算することによる速度向上はかなり大きいです。

同じ測定結果を1秒あたりに処理できるサンプル数に換算して示します。

f:id:select766:20200502093247p:plain
バッチサイズと1秒当たりのサンプル数

16bitの場合で述べますと、極大値がバッチサイズ63、126、252の時に発生しており、そのサイズを超えると一気に性能が悪くなります。2の累乗より少し手前が極大になるようですが、2の累乗-1というわけでもなく実測してみるしかなさそうです。バッチサイズ126のときと252のときでは速度があまり変わらないことを考えると、レイテンシが小さいバッチサイズ126を最大サイズとして用いるのがよさそうです。

追記(2020-05-10): GPUのSM (Streaming Multiprocessor)の数から最適値がある程度予想できるそうです。

これ以降は16bitに固定し、最大バッチサイズ126として最適化プロファイルの効果を確認します。

最適化プロファイルを(1,126,126)の1つだけとした場合と、[(1,1,1),(2,15,15),(16,126,126)]の3つを用いた場合、それと従来の(1,256,256)を比較しました。

f:id:select766:20200504164606p:plain
バッチサイズと計算時間(16bit、最適化プロファイル設定比較)

バッチサイズ126に対して最適化を行うことで、256に対する最適化を行うよりバッチサイズ126時点での速度を高められています。一方、バッチサイズ63付近では劣っています。

f:id:select766:20200504164655p:plain
バッチサイズ20までを拡大

バッチサイズ1,15,126に対して最適化を行うと、まずバッチサイズ1における処理時間を半分ほどに下げることができています。バッチサイズ15までについても、バッチサイズ256に最適化したときと同等の性能となっています。ルート局面付近などの評価ではバッチサイズ1やそれに近い少量のサンプルだけを評価したいため、その遅延を下げられることのメリットがある可能性があります。

TensorRTの最適化はブラックボックスのため試行錯誤するほかないですが、チューニングはかなり影響があるという結果に注意が必要といえます。

*1:n*(1<<14)*SM)/(H*W*C