TensorRTは、NVIDIA社が提供している、Deep Neural Networkの推論を高速に行うライブラリです。NVIDIA社のGPU上での推論(学習済みモデルの実行)に特化しており、CaffeやPyTorchで学習したモデルを読み込んで実行計画を最適化したうえで推論してくれます。もちろんPyTorch等の学習ができるフレームワークでも推論に使えるのですが、推論専用の最適化がなされるところが特徴的で、実験してみたところPyTorchより高速な結果を得ることができました。
私は今回TensorRTを将棋AI向けのDNNの推論に用いました。もともとの動機としてはPythonを用いてPyTorchで学習したモデルをC++で書かれたゲーム木探索部から使いたいという目標があり、C++から使えて推論が高速なライブラリであるとみなすことができるので使いました。PyTorchのモデルをC++から使うためにはLibTorchを使うという手段もあるかと思いますが、推論速度の最適化がデフォルトでなされるので簡単というのがTensorRTのメリットになります。ここでは将棋のDNNを対象にしていますが、出力が2つある点以外は画像分類タスクと同じです。
ソースコードはこちらで公開しています。TensorRTについてはほとんど日本語の記事が見当たらず、サンプルを呼んだものの試行錯誤が必要だった点がいくつかあるためポイントなどを解説します。
モデル形式
TensorRTはいくつかの形式のモデルを読み込むことができますが、PyTorchで学習したモデルを用いる場合はONNX形式が使えます。
DNNはResNetをベースとした畳み込み層主体のもので、画像分類と似ていますが出力が2つあるものを扱います。入力サイズ: (batchsize, 119, 9, 9)、出力サイズ: policy=(batchsize, 2187)、 value=(batchsize, 2)。
ここでのポイントは、バッチサイズを実行時に動的に変えられるようにすることです。そのためにはdynamic_axesというオプションが必要でした。
torch.onnx.export(model, torch.randn(1, 119, 9, 9), "/path/to/output/file", export_params=True, opset_version=10,
verbose=True, do_constant_folding=True, input_names=["input"],
output_names=["output_policy", "output_value"],
# TensorRTでバッチサイズを可変にする際に必要
dynamic_axes={'input': {0: 'batch_size'}, # variable length axes
'output_policy': {0: 'batch_size'},
'output_value': {0: 'batch_size'}})
ここから先はPyTorchはお役御免で、出力されたONNXモデルファイルだけを使います。
環境構築
OS: Ubuntu 18.04を想定します。GPUのドライバ及びCUDA 10.2をインストールしておきます。
cuDNN v7.6.5 (November 18th, 2019), for CUDA 10.2 (cuDNN Library for Linux)およびTensorRT 7.0.0.11 for Ubuntu 18.04 and CUDA 10.2 tar packageを探してダウンロードします。Ubuntu向けのパッケージファイルもあるのですが使い方がよくわかりませんでしたので、tar.gzのほうを使いました。
必要に応じて次のように(雑なやり方ですが)パスを通します。TensorRT自体は実行バイナリ状態で配布されているのでビルド作業はありません。
sudo bash mkdir /usr/local/mycudnn cd /usr/local/mycudnn tar zxvf cudnn-10.2-linux-x64-v7.6.5.32.tgz tar zxvf TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn7.6.tar.gz echo /usr/local/mycudnn/cuda/lib64 >> /etc/ld.so.conf.d/00cudnn.conf echo /usr/local/mycudnn/TensorRT-7.0.0.11/targets/x86_64-linux-gnu/lib >> /etc/ld.so.conf.d/00cudnn.conf ldconfig
TensorRTを使うアプリケーションのコンパイル、リンクについてはいくつかオプションが必要なのでMakefileを参照ください。
大まかな流れ
サンプルコードを抜粋して解説していきます。これは、モデルの推論速度を様々なバッチサイズでテストするベンチマークです。GPUごとにスレッドを立てて同時実行することによりマシン全体でのパフォーマンス測定もできます。
includeしているcommon.hなど、リポジトリに同梱されているものはTensorRT本体ではなく、サンプルコードで使われているユーティリティです。典型的な使い方をする限りはこれらを活用すると楽になります。
TensorRTを使用する際の手順は、
となります。
使用するGPU番号を指定する
if (cudaSetDevice(device) != cudaSuccess)
{
gLogError << "cudaSetDevice failed" << std::endl;
return;
}
cudaSetDeviceで使用するGPUを選択します。0~GPU数-1の整数です。このAPIはTensorRTのものではなくてCUDAを直接叩くことになります。これを呼び出したスレッドで使うGPU番号の指定となります。
ONNXモデルからエンジンをビルドする
auto parsed = parser->parseFromFile(
"data/trt/model.onnx", static_cast<int>(gLogger.getReportableSeverity()));
onnxモデルファイルを読み込みます。
builder->setMaxBatchSize(batchSizeMax);
使用する最大バッチサイズを指定しておく必要があります。
config->setFlag(BuilderFlag::kINT8);
計算を8bitで量子化された状態で行う場合はこのオプションを指定します。ただ、量子化するためのスケール値を適切に指定しないと誤差が大きすぎて使用できません。私が今回使う予定のV100(GPUの型番)では8bit演算コアがないようなので、このオプションは使用していません。
config->setFlag(BuilderFlag::kFP16);
計算を16bit浮動小数点数で行うオプションです。V100ではこれを指定すると(デフォルトの32bitと比べて)倍速以上の速度が出る場合があります。ただし計算誤差が出ます。
最適化プロファイルの作成
TensorRTはモデルとGPUの組み合わせに対して実行計画を最適化しますが、バッチサイズによって適切な実行計画は異なる可能性があります。そのため、複数の実行計画（プロファイル）を使い分ける機能が搭載されています。 プロファイルは、(最小バッチサイズ, 最適バッチサイズ, 最大バッチサイズ)という3つの数値の組を与えて生成します。最適バッチサイズの入力が与えられたときの実行速度が最大となるように実行計画が最適化されます。 小さいバッチサイズに対して、大きいサイズとは別のプロファイルを作成するほうが小さいバッチサイズでの性能が高くなることが期待できます。実験結果については別の記事で書きたいと思いますが、指定方法は以下のようになります。
auto profile = builder->createOptimizationProfile();
profile->setDimensions(inputTensorNames[0].c_str(), OptProfileSelector::kMIN, Dims4{lastbs + 1, 119, 9, 9});
profile->setDimensions(inputTensorNames[0].c_str(), OptProfileSelector::kOPT, Dims4{bs_opt, 119, 9, 9});
profile->setDimensions(inputTensorNames[0].c_str(), OptProfileSelector::kMAX, Dims4{bs_max, 119, 9, 9});
int profileIdx = config->addOptimizationProfile(profile);
for (int b = lastbs + 1; b <= bs_max; b++)
{
profileForBatchSize[b] = profileIdx;
}
OptProfileSelector::kMINでプロファイルが対応する最小バッチサイズ(厳密には、バッチサイズ以外の次元も含めたテンソルの最小サイズ)、OptProfileSelector::kOPTが最適バッチサイズ、OptProfileSelector::kMAXが最大バッチサイズの指定となります。これらを指定したのちconfig->addOptimizationProfileを呼び出すことでプロファイルが登録され、プロファイル番号が得られます。この番号は推論時に必要になりますので、バッチサイズごとにどのプロファイル番号を使用するかを配列に保存しています。
このソースコードでは何やら文字列のパースと絡めていますが、次のような指定ができるようにしています。
profileBatchSizeRange: opt1-max1-opt2-max2...
profileBatchSizeRange=="10-20-100-200"のとき、
- バッチサイズ1~20について、バッチサイズ10に最適化したプロファイルを作成
- バッチサイズ21~200について、バッチサイズ100に最適化したプロファイルを作成
なお、プロファイルは必ずしも複数作る必要はなく、最小バッチサイズ～最大バッチサイズをカバーする1つのプロファイルだけでも十分動作します。
エンジンをビルドした後、推論に必要な「コンテキスト」を作成する必要があります。コンテキストはプロファイルごとに作成する必要があります。1つのコンテキストに対してsetOptimizationProfileを毎回呼び出して対象プロファイルを切り替えるという操作はエラーとなるようです。
for (int i = 0; i < mEngine->getNbOptimizationProfiles(); i++)
{
auto ctx = std::shared_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext(), samplesCommon::InferDeleter());
if (!ctx)
{
return false;
}
ctx->setOptimizationProfile(i);
mContextForProfile[i] = ctx;
}
エンジンのシリアライズ
実行してみると分かりますが、エンジンのビルドには数十秒～数分かかります。ビルド中に様々な実行計画の候補を比較検討しているのだと予想されます。アプリケーション起動時に毎回待たされるのは困るので、ビルド済みのエンジンをファイルに保存して再利用することができます。
IHostMemory *serializedModel = mEngine->serialize(); ofstream serializedModelFile(serializePath, ios::binary); serializedModelFile.write((const char *)serializedModel->data(), serializedModel->size());
シリアライズされたエンジンのロードは以下のように行います。ちょっと煩雑に見えますが、ファイルサイズをチェックした後、単にファイル全体をfdataに読み込んで、runtime->deserializeCudaEngineに渡しているだけです。
ifstream serializedModelFile(serializePath, ios::in | ios::binary); serializedModelFile.seekg(0, ios_base::end); size_t fsize = serializedModelFile.tellg(); serializedModelFile.seekg(0, ios_base::beg); std::vector<char> fdata(fsize); serializedModelFile.read((char *)fdata.data(), fsize); auto runtime = createInferRuntime(gLogger); mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(fdata.data(), fsize, nullptr), samplesCommon::InferDeleter());
なお、エンジンはGPUの機種に固有のものなので、シリアライズしたものを別のマシンに持っていっても動くとは限りません。
推論
CPUとGPUでデータをやり取りするための入出力バッファを作成します。
auto mContext = mContextForProfile.at(profileForBatchSize[batchSize]);
std::string inputBindingName = addProfileSuffix(inputTensorNames[0], profileForBatchSize[batchSize]);
int bidx = mEngine->getBindingIndex(inputBindingName.c_str());
mContext->setBindingDimensions(bidx, Dims4{batchSize, 119, 9, 9});
まずバッチサイズに対応するコンテキストを取り出します。次に、コンテキストにバッチサイズを教えるための手順があります。プロファイルが複数ある場合にこれがトリッキーで、プロファイル0の場合は入力テンソル名(ONNXのエクスポート時に指定した'input')をgetBindingIndexに与えればいいのですが、プロファイル1以降では'input [profile 1]'のようなプロファイル番号を組み合わせた入力テンソル名を与える必要があります。かなり不可解な仕様ですが、これをやると動きます。参考
ここから先はデータのやり取りと実行ですが、特に変なことはありません。samplesCommon::BufferManagerを使えば容易です。
標準出力の抑制
TensorRTは何かとデバッグメッセージが標準出力・標準エラー出力に吐き出されるのですが、将棋AIでは標準入出力を指し手のやり取りに使うので邪魔になりますのでこれを抑制します。
setReportableSeverity(Logger::Severity::kINTERNAL_ERROR);
以上のようなテクニックを駆使することで、将棋AIにTensorRTを組み込むことができます。
