AMPとTensorRT - 48's diary

私個人や関係者にとって自明なことでも意外と記録を残すことは重要かなと思うことがありました。忙しいとサボっちゃうんだけどね。

AMPとTensorRTは昨年から深層学習分野で活躍している技術です。

AMPはAutomatic Mixed Precisionの略で，自動的に浮動小数点の精度を変更して計算速度を上げる技術です。

世間一般では倍精度と単精度の浮動小数点演算が使われています。それぞれ64bit，32bitです。大昔倍精度が普及する前は単精度を組み合わせて精度を上げたりしていましたが和や差でも2倍の時間がかかるのは自明です。

機械学習分野ではこの逆で精度を下げることにより演算速度を上げる技術が今流行っています。特に行列演算では同時に複数の演算を行います。半精度というのが普及しており色々あるのですがfp16やらBF16やらあります。簡単には64bit演算1個する間に16bit演算4個できるだろうといった感じです。

スーパーコンピュータの世界では倍精度の速度が一義ですので，これに対して単精度，半精度を高速化したものが普及しつつありますが，GPUの世界では日常単精度で十分なので倍精度を放置したまま半精度や整数演算が実装されています。

この後者を使った機械学習の高速化のひとつが，PyTorch 1.6から導入されたAMPです。学習モデルの必要な部分に合わせて精度を自動変更してくれるのでとても便利。モデルによりますが2倍近い高速化が実現されています。dlshogi用は昨年夏私が書きましたが，私がやらなくても山岡さんがやってたと思ってます。

dlshogiのautomatic mixed precision対応 - TadaoYamaokaの日記

python-dlshogi2の方にもプルリクを送っておきましたが，こちらは書籍記載事項との兼ね合いもありどういう扱いになるか山岡さんにお任せです。

AMP対応 by bleu48 · Pull Request #4 · TadaoYamaoka/python-dlshogi2 · GitHub

次にTensorRTです。こちらは学習というより推論時に威力を発揮します。

モデル構造を詳細に調べて自動的に高速化するように再配置してくれます。何も考えずに適用しても3割くらい速くなるとかそういう例があります。深層学習のモデルを計算する上での手順を手動で高速化するのは大変ですので使うのがいいでしょう。

TensorRTはNVIDIA社のもので，一般化されたオープンソース実装はtvmというものがあります。といってもNVIDIAのGPU前提の人は一般化する必要はないでしょう。

Apache TVM

さらに，この最適化の際にも半精度化は効きます。ということで，状況次第では数倍の高速化が実現されます。

こちらは上記のAMPより先に山岡さんが昨年4月段階でdlshogiに導入されています。

将棋AIの進捗その42(TensorRT対応) - TadaoYamaokaの日記

二番絞りは上記の高速化技術を前提に層の厚いモデルを作成しました。

それまでは比較的AobaZeroを重点にいじっていたのですが，学習部に再度手を出す動機になったのは以上の高速化によるところが大きいです。