GCP の Tesla T4 が安くなったと思ったら元通りの価格にもどっていて、あれは幻だったのか・・・と嘆いている R&D チームの奥村（@izariuo440）です。これまで何度か TensorRT について触れてきましたが、どのように使うのかは触れていませんでした。今回は、TensorRT の C++ API を使うための主要オブジェクトについて紹介します。Python API については触れません。
TensorRT をなぜ使うのかについてはこちら。
TensorRT の更新ウォッチはこちら。
TensorRT リファレンスはこちら。
概要
TensorRT の利用フェーズは、「コンパイル時」と「実行時」に分けることができます。
コンパイル時には、ビルダーによってネットワーク定義をエンジンに組み込み、エンジンによってホストメモリに内容をシリアライズします。これによってコンパイル済みモデル（ネットワーク）が生成できます。
実行時とは、コンパイル済みモデルを使った推論実行時のことです。ランタイムによってコンパイル済みモデルからエンジンを復元し、エンジンによって実行コンテキストを作って推論を実行します。
コンパイルは、ネットワーク定義にもよりますが数分程度の時間がかかることがあります。アプリケーションで使うときは、事前にコンパイルしてコンパイル済みモデルを作成しておいて、それを読み込んで使うのがよいでしょう。ただし、コンパイル済みモデルは GPU、TensorRT のバージョン、OS などが異なると使えないので注意してください。また、コンパイル時にはプロファイリングによる最適化が行われるので、GPU が何もしていない状態でコンパイルするのがいいでしょう。
コンパイル時と実行時の手順は以下のとおりです。
- コンパイル時
- アプリでビルダーを作る
- ビルダーでビルダー環境設定を作り、必要に応じて変更する
- ビルダーでネットワーク定義を作り、ネットワーク構造を構築する
- ビルダーでネットワーク定義とビルダー環境設定から、エンジンを作る
- エンジンでネットワークをホストメモリにシリアライズする
- アプリがでストメモリにシリアライズされた内容をファイルなどに書き出す
- 実行時
- アプリでランタイムを作る
- アプリでファイルなどからシリアライズされた内容を読み取る
- ランタイムで上記内容をデシリアライズして、エンジンを作る
- エンジンで実行コンテキストを作る
- エンジンで入出力テンソルの情報を取得する
- アプリで上記に応じて CUDA デバイスメモリを確保し、推論用バインディングを用意する
- アプリで CUDA ストリームを作成する
- アプリで入力テンソルを CUDA に転送する
- 実行コンテキストでバインディングと CUDA ストリームを指定して、推論を実行する
- アプリで出力テンソルを CUDA から転送する
- アプリで CUDA ストリームの実行完了を待機する
- アプリで出力テンソルを扱う
ビルダー
nvinfer1::IBuilder クラスが該当します。
ネットワーク定義からエンジンをビルドします。
ビルダー環境設定
nvinfer1::IBuilderConfig クラスが該当します。
ネットワーク定義
nvinfer1::INetworkDefinition が該当します。
ビルダーに入力するネットワーク定義です。
ネットワーク定義は、ネットワーク構造を定義します。ビルダーは、ビルダー環境設定を指定してネットワーク定義をエンジンに組み込みます。ネットワーク定義として、バッチ次元を暗黙的にするか（バッチ次元を実行時に指定）、すべての次元を明示的にするかを選択できます。前者を implicit batch size mode、後者を full dims mode と呼びます。createNetwork() でネットワーク定義を作ると、前者しかサポートされません。hasImplicitBatchSize() でどのモードになっているかを確認できます。
エンジン
nvinfer1::ICudaEngine が該当します。
構築したネットワークで推論を実行するためのエンジンです。機能的に安全でない機能が含まれます。
ホストメモリ
nvinfer1::IHostMemory が該当します。
ライブラリによって確保されたメモリをハンドルするクラスで、それをユーザーがアクセスできるようにします。エンジンによってシリアライズされた内容を、ランタイムがデシリアライズできるように橋渡しするために使います。
ホストメモリオブジェクトによって割り当てられたメモリはライブラリによって所有されていて、destroy メソッドを呼び出すと破棄されます。
ランタイム
nvinfer1::IRuntime が該当します。
シリアライズされた機能的に安全でないエンジンをデシリアライズします。
実行コンテキスト
nvinfer1::IExecutionContext が該当します。
エンジンを使って推論を実行するコンテキストです。機能的に安全でない機能が含まれます。
複数の実行コンテキストが、ひとつのエンジン (nvinfer::ICudaEngine) に対して存在することがあります。これにより、同じエンジンを使って、複数のバッチを同時に実行するような使い方ができます。エンジンが動的 shape をサポートしている場合、同時に使用される各実行コンテキストは、個別に最適化プロファイルを使う必要があります。
TensorRT のプラグイン
TensorRT には深層学習で用いられる多くのオペレーター（レイヤー）が実装されていて、大抵はそれを使えば事足ります。しかし、物体検出などは特殊なオペレーターが必要になることが多く、それらは TensorRT では「プラグイン」という形で実装されています。TensorRT が標準で用意しているプラグインを利用することもできますし、自分で実装することもできます。
プラグインは、TensorRT 4 以前とそれ以降でインターフェースが異なります。
TensorRT 4 以前
IPlugin と IPluginFactory を使う。
TensorRT 5 以降
IPluginV2 と IPluginCreator を使う。IPluginV2 の拡張もある。IPluginV2 は、TensorRT 7 でもまだ非推奨になっていない。
プラグイン毎に以下が必要になる。
IPluginV2の派生クラスを実装するIPluginCreatorの派生クラスを実装するREGISTER_TENSORRT_PLUGINマクロでIPluginCreatorの派生クラスを登録する
IPluginV2
プラグインは、アプリケーションにカスタムレイヤーを実装するためのメカニズムである。IPluginCreator と組み合わせると、プラグインを登録し、デシリアライズするときにプラグインレジストリを検索するするメカニズムを提供する。17 個のインターフェースがあり、うち 16 個のインターフェースを実装する必要がある。
getTensorRTVersion
このプラグインがビルドされた TensorRT API バージョンを返す。プラグインの後方互換性を維持するために TensorRT によって使用されるので、上書きしてはいけない。
virtual int getTensorRTVersion() const TRTNOEXCEPT
{
return NV_TENSORRT_VERSION;
}
getPluginType
プラグインの種類を返す。対応するプラグイン作成器によって返されるプラグイン名 (IPluginCreator::getPluginName()) と一致すべきである。
virtual const char* getPluginType() const TRTNOEXCEPT = 0;
getPluginVersion
プラグインのバージョンを返す。対応するプラグイン作成器によって返されるプラグインバージョン (IPluginCreator::getPluginVersion()) と一致すべきである。
virtual const char* getPluginVersion() const TRTNOEXCEPT = 0;
getNbOutputs
出力テンソルの個数を返す。この関数は INetworkDefinition と IBuilder の実装によって呼び出される。特に、initialize() を呼び出す前に呼び出される。
virtual int getNbOutputs() const TRTNOEXCEPT = 0;
getOutputDimensions
出力テンソルの次元を返す。この関数は INetworkDefinition と IBuilder の実装によって呼び出される。特に、initialize() を呼び出す前に呼び出される。
virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRTNOEXCEPT = 0;
supportsFormat
フォーマットのサポートを確認する。指定された type と format の組み合わせをサポートしている場合、true を返す。この関数は INetworkDefinition と IBuilder と ICudaEngine の実装によって呼び出される。特に、エンジンを生成するときとエンジンをデシリアライズするときに呼び出される。
format 引数について、PluginFormat::kCHW4、PluginFormat::kCHW16、PluginFormat::kCHW32 は渡されることがない。これは TensorRT 5.x シリーズとの後方互換性を維持するためである。他の形式については PluginV2IOExt や PluginV2DynamicExt を使うこと。
virtual bool supportsFormat(DataType type, PluginFormat format) const TRTNOEXCEPT = 0;
configureWithFormat
レイヤーを環境設定する。この関数は initialize() の前にビルダーによって呼び出される。これにより、レイヤーは重み、次元、最大バッチサイズにもとづいてアルゴリズムを選択する機会を提供する。
引き渡される次元は、バッチサイズを含まない。
format 引数について、PluginFormat::kCHW4、PluginFormat::kCHW16、PluginFormat::kCHW32 は渡されることがない。これは TensorRT 5.x シリーズとの後方互換性を維持するためである。他の形式については PluginV2IOExt や PluginV2DynamicExt を使うこと。
virtual void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) TRTNOEXCEPT = 0;
initialize
レイヤーを実行するために初期化する。エンジンの作成時に呼び出される。
成功したら 0、失敗したら 0 でない値を返す（これによりエンジンが終了する）。
virtual int initialize() TRTNOEXCEPT = 0;
terminate
プラグインのレイヤーの初期化中に確保したリソースを開放する。エンジンを破棄したときに呼び出される。
virtual void terminate() TRTNOEXCEPT = 0;
getWorkspaceSize
レイヤーに必要なワークスペースのサイズを見つける。
initialize() の後、エンジンの起動中に呼び出される。返されるワークスペースのサイズは、最大バッチサイズまでのすべてのバッチサイズに対して十分であるべきである。
virtual size_t getWorkspaceSize(int maxBatchSize) const TRTNOEXCEPT = 0;
enqueue
レイヤーを実行する。
成功したら 0、失敗したら 0 でない値を返す（これによりエンジンが終了する）。
virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) TRTNOEXCEPT = 0;
getSerializationState
シリアライズするのに必要なバッファのサイズを見つける。
virtual size_t getSerializationSize() const TRTNOEXCEPT = 0;
serialize
レイヤーをシリアライズする。
buffer は、データをシリアライズするためのバッファへのポインタである。このサイズは getSerializationSize() によって返された値と同じでなければならない。
virtual void serialize(void* buffer) const TRTNOEXCEPT = 0;
destroy
プラグインオブジェクトを破棄する。ネットワーク、ビルダー、エンジンが破棄されるときに呼び出される。
virtual void destroy() TRTNOEXCEPT = 0;
clone
プラグインオブジェクトをクローンする。これにより、内部のプラグインパラメーターがコピーされ、これらのパラメーターを持つ新しいプラグインオブジェクトを返す。
virtual IPluginV2* clone() const TRTNOEXCEPT = 0;
setPluginNamespace
このプラグインが属する名前空間を設定する。理想的には、同じプラグインライブラリのすべてのプラグインオブジェクトは、同じ名前空間を持つべきである。
virtual void setPluginNamespace(const char* pluginNamespace) TRTNOEXCEPT = 0;
getPluginNamespace
このプラグインオブジェクトの名前空間を返す。
virtual const char* getPluginNamespace() const TRTNOEXCEPT = 0;
IPluginCreator
プラグイン生成器。IPluginV2 のプラグインを生成する。
9 個のインターフェースがあり、うち 7 個のインターフェースを実装する必要がある。
getTensorRTVersion
このプラグイン生成器がコンパイルされた TensorRT API バージョンを返す。
virtual int getTensorRTVersion() const TRTNOEXCEPT { return NV_TENSORRT_VERSION; }
getPluginName
プラグインの名前を返す。
virtual const char* getPluginName() const TRTNOEXCEPT = 0;
getPluginVersion
プラグインのバージョンを返す。
virtual const char* getPluginVersion() const TRTNOEXCEPT = 0;
getFiledNames
createPlugin に引き渡される必要のあるフィールド名のリストを返す。
virtual const PluginFieldCollection* getFieldNames() TRTNOEXCEPT = 0;
createPlugin
プラグインオブジェクトを返す。エラー時には nullptr を返す。
virtual IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) TRTNOEXCEPT = 0;
deserializePlugin
プラグインレイヤーのデシリアライズ中に呼び出される。プラグインオブジェクトを返す。
virtual IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) TRTNOEXCEPT = 0;
setPluginNamespace
属するプラグインライブラリにもとづいてプラグイン生成器の名前空間を設定する。プラグイン生成器の登録中 (IPluginRegistry::registerCreator()) に設定できる。
virtual void setPluginNamespace(const char* pluginNamespace) TRTNOEXCEPT = 0;
getPluginNamespace
プラグイン生成器オブジェクトの名前空間を返す。
virtual const char* getPluginNamespace() const TRTNOEXCEPT = 0;
~IPluginCreator
デストラクタ。
virtual ~IPluginCreator() {}
IPluginRegistry
アプリケーション内のすべてのプラグインの単一登録所。エンジンのデシリアライズ中にプラグインの実装を見つけるために使われる。内部的には、プラグインレジストリはシングルトンとみなされるため、アプリケーション内のすべてのプラグインは、同じグローバルなレジストリの一部となる。プラグインレジストリは、IPluginV2 のプラグインのみをサポートしており、それに対応する IPluginCreator の実装も必要であることに注意。
REGISTER_TENSORRT_PLUGIN マクロによって、暗黙的に使用される。このマクロ以外でアクセスする機会はあまりないだろう。
まとめ
ソフトウェアの情報は英語がほとんどなのは世の常ですが、日本語情報が少ないのでざっくりとまとめてみました。
オプティムでは、こうした技術に興味がある・作ってみたい・既に作っている、というエンジニアを募集しています。興味のある方は、こちらをご覧ください。
