オプティムの R&D チームで Deep な画像解析をやっている奥村です。手元に NVIDIA RTX A4000 が届いてわくわくしています（PCIe 3.0 でも動作する模様）。TensorRT 8 の変更点についてメモしました。

はじめに

TensorRT 8.0.1 GA は、TensorRT 7.2.3 のリリースから約四ヶ月ぶりのリリースとなります。今回は５月に TensorRT 8.0.0 EA のリリースを挟んでからの GA となっています。気になった内容がいくつかあったので、TensorRT 8.0.1 のリリースノートをベースに、意訳（直訳）したものをメモしました。公式情報ではないので参考程度にしてください。TensorRT については特にここでは触れませんが、 TensorRT で物体検出・姿勢推定はどれくらい速くなる？ - OPTiM TECH BLOG、TensorRT 6 でさらに快適な高速推論 - OPTiM TECH BLOG、TensorRT 7 でさらに快適な高速推論 - OPTiM TECH BLOG で少し触れていますので興味があれば参照してみてください。

気になった内容

QAT で訓練されたネットワークのサポート

INT8 量子化の際に、QAT (Quantization-Aware Training) で訓練されたネットワークがサポートされました。INT8 量子化には、推論速度の向上・使用メモリの削減・エンジンファイルサイズの削減などの効果がありますが、これまでは Post-Training Quantization にしか対応していなかったため精度が少し落ちてしまうという課題がありました。QAT で訓練されたネットワークを使うことにより、精度劣化を抑えることができます。PyTorch 本家の QAT、TensorFlow 本家の QAT、NVIDIA の pytorch-quantization による QAT などで訓練し、それを ONNX モデルを出力してから TensorRT に読み込ませることができます。

精度や速度については以下が参考になります。

developer.nvidia.com

リサイズ方法の強化

リサイズは、深層学習フレームワークごとに実装が異なっており、TensorRT もそのひとつでしたが、TensorFlow・PyTorch・ONNX 由来のリサイズに対応しました。これで微妙に結果が異なってしまうという減少が軽減されると思います。

NMS の試験的なサポート

ONNX の NonMaxSuppression 演算が試験的にサポートされました。出力テンソルの形状が動的ではなく固定であるなどの制限はあるものの、よいですね。

Release ONNX-TensorRT 8.0 Release · onnx/onnx-tensorrt · GitHub

EfficientDet サンプルの追加

TensorFlow な EfficientDet を ONNX に変換してから TensorRT で推論させるサンプルが追加されました。

github.com

巷では YOLOX が話題ですが、EfficientDet もいいですよね。

物体検出器 EfficientDet をブラウザで高速推論 - OPTiM TECH BLOG

意訳

今回のリリースノートはボリュームがあったため「制約事項」「非推奨機能と削除された機能」「修正された問題」「お知らせ」「既知の問題」については割愛させていただきました。

TensorRT 8.0.1 GA

これは TensorRT 8.0.1 のリリースノートです。x86 の Linux と Windows、PowerPC/ARM Server Base System Architecture (SBSA) の Linux ユーザーにのみ適用されます。

このリリースノートは、特に記載のない限りワークステーション・サーバー・JetPack ユーザーに適用されます。Jetson プラットフォームには適用されませんん。

このリリースは、以前の TensorRT 8.x.x リリースからの修正と以下の変更を含みます。以前リリースされたバージョンについては TensorRT Documentation Archives を参照してください。

主要な機能と改善

RedHat/CentOS 8.3、Ubuntu 20.04、SUSE Linux Enterprise Server 15 Linux ディストリビューションがサポートされました。現時点では SLES 15 向けは tar ファイルによるインストールのみです。詳細は TensorRT Installation Guide を参照してください。
Python 3.9 がサポートされました。
IResizeLayer に、ResizeCoordinateTransformation・ResizeSelector・ResizeRoundMode という3つの列挙型が追加され、TensorFlow・PyTorch・ONNX 由来のリサイズモードのサポートが強化されました。詳細は IResizeLayer セクションを参照してください。
Builder Layer Timing Cache がシリアライズ可能になり、ビルダーのインスタンスをまたいで再利用できるようになりました。詳細は Builder Layer Timing Cache や trtexec を参照してください。
カーネルの重みが Structured Sparsity である畳み込みと全結合を追加しました。この機能は IBuilderConfig の kSPARSE_WEIGHTS フラグを設定することで有効になりますが、NVIDIA Ampere 世代の GPU でしか利用できません。詳細は Structured Sparsity を参照してください。
２つの新しいレイヤーが追加されました。IQuantizeLayer と IDequantizeLayer は、演算とデータバッファの精度を明示するために仕様できます。ONNX の QuantizeLinear と DequantizeLinear 演算はこれらのレイヤーにマップされ、Quantization-Aware Training (QAT) で訓練されたネットワークのサポートを可能にします。詳細は、Explicit-Quantization、IQuantizeLayer、IDequantizeLayer、Q/DQ Fusion を参照してください。
1D fused depthwise + pointwise 畳み込みカーネルのサポートによって QuartzNet の最適化を行い、A100 で 1.8 倍の性能改善を達成しました（Jetson プラットフォームには適用されません）。
ONNX 演算の Celu、CumSum、EyeLike、GatherElements、GlobalLpPool、GreaterOrEqual、LessOrEqual、LpNormalization、LpPool、ReverseSequence、 SoftmaxCrossEntropyLoss をサポートしました。詳細は Supported Ops や Supported ONNX Operators を参照してください。
DLA で INT8 の Sigmoid/Tanh をサポートしました。これにより、Sigmoid/Tanh を含む DLA サブグラフを、内部で FP16 に自動アップグレードすることにより、INT8 にコンパイルできるようになります。詳細は DLA Supported Layers を参照してください。
DLA で native planar format と native gray-scale format をサポートしました。
EngineCapability が EngineCapability::kDEFAULT のとき、DLA で再フォーマット不要なエンジンを生成できるようになりました。
TensorRT は、例外がライブラリの境界を超えてはいけないことを明示するために noexpect キーワードを使用して API を宣言するようになりました。IGpuAllocator や IPluginV2 のようにアプリケーションが継承するすべての TensorRT クラスは、TensorRT によって呼び出されたメソッドが補足されない例外を送出しないことを保証する必要があります。そうしないと動作が未定義になります。
TensorRT は、すべてのエラーを ErrorRecorder API を介して関連する ErrorCode とともにエラーを報告します。ErrorRecorder は、エラーレコーダーが登録されていない場合、Severity::kERROR または Severity::kINTERNAL_ERROR でレガシーなロガーにフォールバックします。ErrorCodes を使うと、TensorRT が以前に回復不可能な状況を報告した場合に回復できます。
EfficientNet などいくつかの CNN で仕様される GlobalAveragePooling 演算の性能が向上しました。INT8 精度の Transformer ベースのネットワークの場合、Quantization Aware Training (QAT) で訓練したネットワークを使い、ネットワーク定義に IQuantizeLayer と IDequantizeLayer を使うことをお勧めします。
名前つき再フィット重みをサポートしました。詳細は Refitting An Engine を参照してください。
再フィットの性能が向上しました。性能向上は、重みが大きい場合や、重みやレイヤーがいくつも同時に更新される場合に明らかになります。
以下のサンプルを追加しました。
1. engine_refit_onnx_bidaf: ONNX BiDAF モデルからエンジンを構築し、再フィットする。詳細は Refitting An Engine Built From An ONNX Model In Python を参照してください。
2. efficientdet: Google EfficientDet を変換・実行する。詳細は Scalable And Efficient Object Detection With EfficientDet Networks In Python を参照してください。
BERT のような Transformer ベースのネットワークや、Multi-Head Self-Attention を使用するネットワークのパフォーマンスが向上しました。
IBuilderConfig::setTacticSources に指定できる列挙子として cuDNN が追加されました。演算実装のソースとして cuDNN を使うには、IBuilderConfig::setTacticSources API を使用して有効・無効にできます。
以下の C++ API が追加されました。
1. class IDequantizeLayer
2. class IQuantizeLayer
3. class ITimingCache
4. IBuilder::buildSerializedNetwork()
5. IBuilderConfig::getTimingCache()
6. IBuilderConfig::setTimingCache()
7. IGpuAllocator::reallocate()
8. INetworkDefinition::addDequantize()
9. INetworkDefinition::addQuantize()
10. INetworkDefinition::setWeightsName()
11. IPluginRegistry::deregisterCreator()
12. IRefitter::getMissingWeights()
13. IRefitter::getAllWeights()
14. IRefitter::setNamedWeights()
15. IResizeLayer::getCoordinateTransformation()
16. IResizeLayer::getNearestRounding()
17. IResizeLayer::getSelectorForSinglePixel()
18. IResizeLayer::setCoordinateTransformation()
19. IResizeLayer::setNearestRounding()
20. IResizeLayer::setSelectorForSinglePixel()
21. IScaleLayer::setChannelAxis()
22. enum ResizeCoordinateTransformation
23. enum ResizeMode
24. BuilderFlag::kSPARSE_WEIGHTS
25. TacticSource::kCUDNN
26. TensorFormat::kDLA_HWC4
27. TensorFormat::kDLA_LINEAR
28. TensorFormat::kHWC16
以下の Python API が追加されました。
1. class IDequantizeLayer
2. class IQuantizeLayer
3. class ITimingCache
4. Builder.build_serialized_network()
5. IBuilderConfig.get_timing_cache()
6. IBuilderConfig.set_timing_cache()
7. IGpuAllocator.reallocate()
8. INetworkDefinition.add_dequantize()
9. INetworkDefinition.add_quantize()
10. INetworkDefinition.set_weights_name()
11. IPluginRegistry.deregister_creator()
12. Refitter.get_all_weights()
13. Refitter.get_missing_weights()
14. Refitter::set_named_weights()
15. IResizeLayer.coordinate_transformation
16. IResizeLayer.nearest_rounding
17. IResizeLayer.selector_for_single_pixel
18. IScaleLayer.channel_axis
19. enum ResizeCoordinateTransformationDoc
20. enum ResizeMode
21. BuilderFlag.SPARSE_WEIGHTS
22. TacticSource.CUDNN
23. TensorFormat.DLA_HWC4
24. TensorFormat.DLA_LINEAR
25. TensorFormat.HWC16
Linux プラットフォーム上でメモリレポーティングメカニズムが改善されました。以前は追跡されていなかった大きな割当が含まれるようになりました（ヒープ割当ではなくメモリマップしていなことが原因）。

破壊的な API 変更

Python 2 のサポートを停止しました。つまり、TensorRT に Python 2 の wheel が含まれなくなり、Python のサンプルが Python 2 で機能しなくなります。
すべての API が適切に noexcept としてマークされました。IErrorRecorder インターフェースは、エラー報告のための API として完全に統合されました。Logger は、ErrorRecorder がユーザーによって提供されなかった場合のフォールバックとしてのみ使用されます。
コールバックは noexcept としてマークされました。従って、実装も noexcept としてマークする必要があります。TensorRT は、コールバックによって送出された例外を取り扱うことはありませんでしたが、API によって補足されるようになりました。
ポインタ配列が変更できないことを示す void** 型のパラメータを受け取るメソッドが、void *const * 型を受け取るように変更されました。
Dims は Dims32 クラスの型愛リアスになりました。Dims を前方宣言するコードは、class Dims32; using Dims = Dims32; と変更すべきです。
TensorRT 8.0 EA から TensorRT 8.0 GA への変更で、getLogger() 関数プロトタイプが NvInferRuntimeConmmon.h から NvInferRuntime.h に移動されました。getLogger() を使用していて NvInferRuntimeCommon.h のみを include していたコードは、変更する必要があるかもしれません。

互換性

TensorRT 8.0.1 は、以下でテストされています。
この TensorRT リリースでは以下の CUDA をサポートしています。
テスト済みのソフトウェアスタックで TensorRT を使用することを推奨します。ソフトウェアスタックには、Features For Platforms And Software セクションに記載されている cuDNN と cuBLAS が含まれます。 cuDNN と cuBLAS の他の互換リリースを使用することもできますが、他のバージョンでは性能が退行することがあります。まれに、機能の退行が見られることがあります。

まとめ

少し前に ONNX 2020 にも記載しましたが、ONNX 対応が強化されてかなりのモデルがプラグインなしで動くようになってきています。あとは GatherND に対応してくれたり、GatherElements が動的形状を持つ入力に対応してくれたりするとかなり助かりますね。

cuDNN はじめバイナリがどんどん大きくなっているので、静的リンクしたり、nvprune で不要なアーキテクチャ向けのコードを削るなどしてバイナリを小さくできるとうれしいのですが、FP16 推論しようとすると落ちるという問題に遭遇しています。この辺もいずれ解消されると嬉しいですね。

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