TensorRT 歴３年となった R&D チームの奥村（@izariuo440）です。今回は、2021/04/12 に発表された Ampere 世代の各種 GPU に対して TensorRT でその性能を十分に引き出すために必要な知識をまとめておきます。FP16/INT8 をうまく使うと、推論速度が２〜４倍になったりするのでおすすめです。

演算精度

NVIDIA TensorRT は、NVIDIA GPU 上で深層学習モデルを高速に推論するためのプラットフォームです。TensorRT を使うモチベーションについては以下をご覧ください。

TensorRT で物体検出・姿勢推定はどれくらい速くなる？ - OPTiM TECH BLOG

TensorRT では推論時の演算精度を選択できます。どの GPU も、32 ビット単精度浮動小数点数である FP32 という演算精度をサポートしています。GPU によっては、16 ビット半精度浮動小数点数である FP16、8 ビット整数である INT8 という演算精度をサポートしています。多くの場合、FP16 や INT8 を選択すると、FP32 よりもさらに高速に推論できます。

演算精度 FP16/INT8 をサポートしている GPU

前述したとおり、FP16/INT8 はすべての GPU がサポートしている訳ではありません。どの GPU がサポートしているかは CUDA Compute Capability によって決定されます。TensorRT Support Matrix の Hardware and Precision には、CUDA Compute Capability ごとにどういう演算精度が利用できるかが表としてまとめられています。

FP16/INT8 に注目して表を整理し、GPU を入れるとこのようになります。

CCC	FP16	INT8	FP16'	INT8'	Jetson	TITAN	GeForce	Quadro	Tesla
8.6	✓	✓	✓	✓			GeForce RTX 3050~	Quadro RTX A4000~	A10
8.0	✓	✓	✓	✓					A100
7.5	✓	✓	✓	✓		NVIDIA TITAN RTX	GeForce RTX 2060~	Quadro RTX 4000~	Tesla T4
7.2	✓	✓	✓	✓	Jetson Xavier
7.0	✓	✓	✓			NVIDIA TITAN V		Quadro GV100	Tesla V100
6.2	✓				Jetson TX2
6.1		✓				NVIDIA TITAN X~	GeForce GTX 1050~	Quadro P400~	Tesla P4~
6.0	✓							Quadro GP100	Tesla P100
5.3	✓				Jetson TX1/Nano

CCC: CUDA Compute Capability の略記
FP16': FP16 Tensor Cores の略記
INT8': INT8 Tensor Cores の略記

CCC 7.0 以前、Pascal 世代や多くの Volta 世代の GPU では、FP16/INT8 のサポート状況が入り組んでいましたが、CCC 7.2 以降というか Turing 世代以降は、FP16/INT8 をサポートしており、対応する Tensor Core も搭載しています。そのため、TensorRT で FP16/INT8 といった演算精度を指定すると、FP32 よりも高速に推論できます。

CCC 8.6 の行の Quadro/Tesla については、私の憶測ですのでご注意ください。

FP16 の威力

TensorRT で FP16 が威力を発揮するのは、上表で FP16 にチェックが入っている GPU だと思われます。TensorRT C++ API だと nvinfer1::IBuilder::platformHasFastFp16()メソッドで確認できます。有効にするには、エンジンを作成するときに nvinfer1::IBuilderConfig::setFlag(nvinfer1::BuilderFlags:: kFP16) を使います。TensorRT のドキュメントにあるように基本的には以下のようなコードを追加するだけです。

nvinfer1::IBuilderConfig *config = builder->createBuilderConfig();
...
config->setFlag(nvinfer1::BuilderFlags:: kFP16);
...
nvinfer1::ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config);
...

FP16 のみチェックが入っているのは、Jetson TX1/TX2/Nano などがあります。この場合、推論速度を FPS で比較すると、FP32 の約２倍になるそうです*1。

FP16 と FP16' の双方にチェックが入っているのは、Volta 世代以降の GPU が該当しています。この場合、推論速度を FPS で比較すると、FP32 の２倍以上になることがあるようです。Jetson Xavier では２倍、Quadro RTX 4000 では３倍になりました *2。また、GPU のメモリ使用量が FP32 の 1/2 程度に減ることも確認しています。

とにかく手軽に高速になるので、使わない手はありませんね。ただし、若干精度が劣化する可能性があるので、十分確認した上で利用するようにしましょう。

INT8 の威力

TensorRT で INT8 が威力を発揮するのは、上表で INT8 にチェックが入っている GPU です。TensorRT C++ API だと nvinfer1::IBuilder::platformHasFastInt8() メソッドで確認できます。有効にするには、エンジンを作成するときに nvinfer1::IBuilderConfig::setFlag(nvinfer1::BuilderFlags::kINT8) を使い、キャリブレーションを実施する必要があります。

推論速度を FPS で比較すると、FP32 の２〜４倍程度になることがあるようです。GeForce GTX 1080/GeForce GTX 1080 Ti では２倍、Quadro RTX 4000/Tesla T4/Jetson Xavier NX では４倍になりました *3。また、GPU のメモリ使用量が FP32 の 1/4 程度に減ることも確認しています。FP16 と同様に、精度が劣化する可能性があるので、十分確認した上で利用するようにしましょう。とはいえ、画像分類ではコンマ数％しか精度が劣化しない、という情報もあるようです。

8-bit Inference with TensorRT

TensorRT のバージョン

TensorRT は CUDA、cuDNN に依存しています。それぞれどのバージョンを使うのか、組み合わせを検討する必要があります。TensorRT のリリースノートには、どういう組み合わせでテストされているのか記載されているので、それを足がかりにするとよいのかもしれません。

TensorRT 7.2.3 のリリースノートには以下のような記載があります。

This is the TensorRT 7.2.3 GA release notes for Windows and Linux x86 users. For NVIDIA Jetson Linux for Tegra users, TensorRT 7.2.3 is an Early Access (EA) release specifically for MLPerf Inference. For production use of TensorRT, we recommend using the TensorRT 7.1.3 GA.

最後の一文にプロダクションで使うなら TensorRT 7.1.3 GA の利用を推奨、とありますが、 Jetson に限った話なのかどうか解釈が難しい・・・。

Ver		Date	CUDA 9	CUDA 10	CUDA 11	cuDNN	ONNX
7.2.3	GA	2021/02/23		2	0, 1, 2	8.1.1	1.6.0
7.2.2		2020/12/11		2	0, 1, 2	8.0.5	1.6.0
7.2.1		2020/10/07		2	0, 1	8.0.4	1.6.0
7.2.0	RC	2020/08/21		2	0	8.0.2	1.6.0
7.1.3	GA	2020/06/17		2	0	8.0.0	1.6.0
7.1.2	RC	?		2	0	8.0.0	1.6.0
7.1.0	EA	?		2		8.0.0	1.6.0
7.0.0		2019/12/17	0	0, 2		7.6.5	1.6.0
6.0.1		2019/09/14	0	0, 1, 2		7.6.5	1.5.0
5.1.5		2019/04/27	0	0, 1		7.5.0	1.4.1

日付は、tarball を展開してファイルの変更日を記載
CUDA 列は、マイナーバージョンのみを記載

互換性

シリアライズした TensorRT エンジンはファイルなどに保存することができます (*.trt など)。しかし、異なる TensorRT バージョン間では原則として互換性がないので注意してください。また、ある GPU で作成した TensorRT エンジンは、他の GPU では原則として互換性がありませんので注意してください。

さいごに

NVIDIA A10 は NVIDIA Tesla T4 の後継となる GPU のようで、推論性能は Tesla T4 の 2.5 倍になるようです。AWS や GCP などで利用可能になるのが待ち遠しいですね。

blogs.nvidia.com

オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。

www.optim.co.jp

*1:噂を聞いただけで、私は動作確認できていません。

*2:速度については、推論するモデルによっても変動しますので参考程度ということでお願いします。

*3:速度については、推論するモデルによっても変動しますので参考程度ということでお願いします。