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

遠藤です。

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

• 第1回: TensorRT の概要について
• 第2回: インストール方法について
• 第4回: 性能検証レポート

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

基本的な流れ

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

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

実際の使い方

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

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

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

#include <NvInfer.h>         // TensorRT 本体
#include <NvCaffeParser.h>   // Caffe モデルのパーサ

(1): TensorRT の初期化

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

// Logger for TensorRT
class Logger : public ILogger {
public:
    void log(ILogger::Severity severity, const char* msg) override
    {
        // suppress information level log
        if (severity == Severity::kINFO) return;
        std::cout << msg << std::endl;
    }
};

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

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

IBuilder* builder = createInferBuilder(logger_);

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

IRuntime* infer = createInferRuntime(logger_);

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

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

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

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

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

    // parser caffe model
    INetworkDefinition* network = builder->createNetwork();
    ICaffeParser* parser = createCaffeParser();

    const IBlobNameToTensor* blobNameToTensor = parser->parse(model_file.c_str(),
                                                              trained_file.c_str(),
                                                              *network,
                                                              DataType::kFLOAT);

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

    // mark output of network (caffe model doesn't have output info)
    for (auto& s : outputs) {
        std::cout << "[INFO] marking blob " << s << " as output." << std::endl;
        CHECK(blobNameToTensor->find(s.c_str()));    // 対応する blob がないと、ここが NULL になる
        network->markOutput(*blobNameToTensor->find(s.c_str()));
    }

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

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

    // build TensorRT engine
    builder->setMaxBatchSize(g_batch_size);
    builder->setMaxWorkspaceSize(16 << 20);

    ICudaEngine* engine = builder->buildCudaEngine(*network);

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

    IHostMemory* modelStream = engine->serialize();

(4): 推論実行

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

    IExecutionContext* context;
    context = engine->createExecutionContext();

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

    // 非同期実行のため、ストリームを生成
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    // 入力データの Host to Device 転送
    for (int i=0; i<num_inputs(); ++i) {
        Blob<Dtype>* b = &blobs[inputIdx_[i]];
        cudaMemcpyAsync(b->gpu_data(), (void *)b->cpu_data(),
                        b->nBytes(), cudaMemcpyHostToDevice, stream);
    }

    // 入出力データの GPU アドレス一覧を作成し、推論を実行
    std::vector<void*> buffers(blobs_.size());
    for (int i=0; i<blobs_.size(); ++i) {
        buffers[i] = blobs[i].gpu_data();
    }
    context->enqueue(g_batch_size, &buffers[0], stream, nullptr);

    // 出力データの Device to Host 転送
    for (int i=0; i<num_outputs(); ++i) {
        Blob<Dtype>* b = &blobs[outputIdx_[i]];
        cudaMemcpyAsync((void *)b->cpu_data(), b->gpu_data(),
                        b->nBytes(), cudaMemcpyDeviceToHost, stream);
    }

    // ストリームの実行完了待ち
    cudaStreamSynchronize(stream);

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

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

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

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

    const IBlobNameToTensor* blobNameToTensor = parser->parse(model_file.c_str(),
                                                              trained_file.c_str(),
                                                              *network,
                                                              DataType::kHALF);  // ココが変わる

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

    builder->setMaxBatchSize(g_batch_size);
    builder->setMaxWorkspaceSize(16 << 20);
    builder->setHalf2Mode(true);    // ココが追加

    ICudaEngine* engine = builder->buildCudaEngine(*network);

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

参考資料一覧

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

• TensorRT Developer Guide
• TensorRT C++ ドキュメント
• TensorRT Python ドキュメント
• TensorRT サンプルコード (TensorRT をインストールしたマシンの /usr/src/tensorrt/samples  に展開されています)

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

このサイトはスパムを低減するために Akismet を使っています。コメントデータの処理方法の詳細はこちらをご覧ください。