【TensorRT やってみた】(3): TensorRT の使い方 (C++ API)

遠藤です。

TensorRT やってみたシリーズの第3回です。

今回は、TensorRT を C++ から呼び出す方法を解説します。TensorRT は API のドキュメント等があまり十分ではないため、参考になると幸いです。

基本的な流れ

TensorRT を利用する際は、以下のステップを踏んでいきます。

  1. TensorRT の初期化
  2. ネットワークモデルの読み込み
  3. 推論エンジンのビルド
  4. 推論実行

実際の使い方

それでは、具体的な使い方をコード例を交えて解説します。

(0): TensorRT のインクルード

意外にもデベロッパーガイド等の公式ドキュメントには、どのヘッダファイルを読み込めばいいのかが書かれていません。TensorRT を利用するには、以下の通りヘッダをインクルードします。

(1): TensorRT の初期化

TensorRT は、内部でエラーメッセージやログ情報を生成しますが、それらをどのように出力するかはユーザに委ねられています。具体的には、下記のような log 関数を持つクラスのインスタンスを生成し、その参照を初期化時に渡します。

TensorRT のワークフローでは、ネットワークモデルの読み込み・推論エンジンビルドと推論の実行は別々になっていて、それぞれ初期化をよびだします。

ネットワークモデルの読み込みと推論エンジンのビルドを行う際には、以下の初期化関数を呼びます。

推論を実行する際には、以下の初期化関数を呼びます。

(2): ネットワークモデルの読み込み

ネットワークモデルの読み込み方法には、以下の3種類があります。

  1. Caffe モデルの読み込み
  2. UFF(Universal Framework Format) モデルの読み込み
  3. C++ API を利用して自分でネットワークを定義、重みファイルを読み込み

この記事では、Caffe モデルの読み込み方法を説明します。UFFモデルの読み込みはサンプルコードの  sampleUffMNIST  を、C++ API を利用した読み込みはサンプルコードの  sampleMNISTAPI  を参照してください。

まず、パーサのインスタンスを生成した後、Caffe のモデルファイルを読み込みます。Caffe のネットワーク定義ファイル ( deploy.prototxt)  と 学習したネットワークのファイル ( *.caffemodel) のファイルのパスをパーサに渡します。

Caffe のネットワーク定義ファイルには、入力レイヤの情報はあるものの出力レイヤの情報がありません。一方で、 TensorRT は入力と出力をそれぞれ必要とします。そこで、出力レイヤの情報を TensorRT に次のように伝えます。

(3): 推論エンジンのビルド

読み込んだネットワークから、推論エンジンをビルドします。コード例は以下のとおりです。

ここで生成された推論エンジンは、以下のコード例のようにシリアライズすることができます。シリアライズした推論エンジンをファイルなどに保存したりすることもできて、便利です。

(4): 推論実行

推論を実行するには、まず推論実行コンテキストを生成します。

次に、推論を実行していきます。コード例は以下のとおりです。

Appendix: 半精度浮動小数点数 (FP16) を使う

FP16 を使う際には、上記の使い方から以下の項目を若干変える必要があります。

  • ネットワークモデルの読み込み
  • 推論エンジンのビルド

ネットワークモデルの読み込みは、次のようになります。

また、推論エンジンのビルドのところも、次のように若干変更となります。

このように、1行の追加と1行の変更だけで FP16 による推論が簡単に実現できるところがポイントです。

参考資料一覧

実際にコードを書く際に参考になる参考資料を一覧にまとめました。特にサンプルコードは具体的にどのようにコードを書けばいいかのいい参考になりますので、TensorRT を使う際にはぜひ参考にしてください。

次回は、実際に TensorRT を呼び出してみて、その性能を見ていきたいと思います。

コメントを残す

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です

This site uses Akismet to reduce spam. Learn how your comment data is processed.