TensorRTの構築

最も簡単な作成および操作プロセスを提供します。

手順

グローバルオブジェクトを作成するILogger

ILoggerこれは抽象クラスであり、使用するには派生させる必要があります。例は公式APIに記載されています。特定の使用法は、コードサンプルからコピーできます。

タイプのオブジェクトを作成する IBuilder

IBuilder* builder = createInferBuilder(Logger);

タイプのオブジェクトを作成する IBuilderConfig

IBuilderConfig* config = builder->createBuilderConfig();

タイプのオブジェクトを作成する INetworkDefinition

// default creation
INetworkDefinition* network = builder->createNetworkV2(0U);

// explicit batch
INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));

ネットワーク定義
エンジンの作成

ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

実行コンテキストの作成

IExecutionContext* context = engine->createExecutionContext();

コンテキストで推論を行う

以下の例は非同期です。（非同期実行）

  // Pointer to Input & Output Buffers
  void* buffers[2];
  
  // In order to bind the buffers, we need to know the names of the input and output tensors.
  // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
  const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);

  // Create GPU buffers on device
  CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
  CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));

  // Create stream
  cudaStream_t stream;
  CUDA_CHECK(cudaStreamCreate(&stream));

  // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
  context->enqueue(batchSize, buffers, stream, nullptr);
  CUDA_CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
  cudaStreamSynchronize(stream);

  // Release stream and buffers
  cudaStreamDestroy(stream);
  CUDA_CHECK(cudaFree(buffers[inputIndex]));
  CUDA_CHECK(cudaFree(buffers[outputIndex]));

簡単に言うと、データのコピーDevice、context->enqueue()実行、コピーバックの3つのステップがありますHost。

ネットワーク定義

// set the input
auto data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{-1, 1, INPUT_H, INPUT_W});

// add layers
auto conv1 = network->addConvolution(*data->getOutput(0), 20, DimsHW{5, 5}, weightMap["conv1filter"], weightMap["conv1bias"]);
conv1->setStride(DimsHW{1, 1});

auto pool1 = network->addPooling(*conv1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2});
pool1->setStride(DimsHW{2, 2});

auto ip1 = network->addFullyConnected(*pool1->getOutput(0), 500, weightMap["ip1filter"], weightMap["ip1bias"]);
auto relu1 = network->addActivation(*ip1->getOutput(0), ActivationType::kRELU);

auto prob = network->addSoftMax(*relu1->getOutput(0));
prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);

// set output
network->markOutput(*prob->getOutput(0));

入力の次元はすべてでDims3あり、暗黙的に指定できますbatchsize。完全に接続されたレイヤーの場合、入力が{C、H、W}の場合、{1、C * H * W}に変換されてから、ネットワークに入力されます。
入力と出力はネットワークITensorベースです。意志ITensorと文字列のバインド、後で使用するためのインデックス。

本文中のコードはここで参照されています

TensorRT-建設

記事のディレクトリ

TensorRTの構築

手順

ネットワーク定義

おすすめ