TensorRT學習（二）通過C++使用

  本文源於學習TensorRT文檔《TensorRT-Developer-Guide》第2章“WORKING WITH TENSORRT USING THE C++ API”的理解。

一、TensorRT實例化對象

  使用TensorRT進行推理需要創建的對象：

• IExecutionContext，用於推理的對象，通過ICudaEngine對象獲取；
• ICudaEngine，TensorRT的引擎，通過對模型序列化buildCudaEngine(*network)、讀取已經序列化的文件deserializeCudaEngine()兩種方式產生，後者在前者基礎上通過serialize()獲得，因此使用更方便；
• ILogger，全局變量，在大多數TensorRT的API中都會作爲參數使用；
• IBuilder，通過 createInferBuilder(gLogger) 創建；
• iNetworkDefinition，在IBuilder對象中調用createNetwork()獲取；
• iParser，以INetwork作爲輸入來創建的網絡解釋器，其內部的parse()方法可以讀取模型文件，這裏要注意的是在Caffe、Uff、ONNX中使用時不同的，具體可看TensorRT的相關示例；
• IRuntime，通過 createInferRuntime(gLogger) 創建。

  ICudaEngine可以通過IBuilder的buildCudaEngine()、IRuntime的deserializeCudaEngine()產生。CUDA的context建議在IRuntime、IBuilder前創建並配置，因爲IRuntime和IBuilder在創建時會使用到context，在context爲手動創建時會使用默認的context。

二、創建網絡

  使用TensorRT進行推理需要創建TensorRT的網絡，有兩種方法：

• 通過iParser序列化已有模型，主要針對Caffe、ONNX、Uff（TensorFlow）;
• 通過TensorRT的API直接定義每一層網絡結構，並填充自己訓練的參數。

  無論是哪種方法，都需要指定輸入輸出節點並命名，非輸出節點會被優化。網絡的權重和模型的優化會被保存在IBuilder的ICudaEngine中，當網絡模型是通過iParser創建時，iParser會有內存管理權，因此iParser的釋放應在IBuilder運行結束之後。

2.1 通過API定義網絡

  sampleMNISTAPI中有關於手動創建每層網絡的示例，其過程爲：

// 定義IBuilder和iNetworkDefinition
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

// 添加輸入層
auto data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{1, INPUT_H, INPUT_W});

// 添加捲積層（輸入節點、strides、權重、偏置）
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);

// 添加Softmax層計算最後的概率並設置節點名稱
auto prob = network->addSoftMax(*relu1->getOutput(0));
prob->getOutput(0)->setName(OUTPUT_BLOB_NAME);

// 設定輸出層
network->markOutput(*prob->getOutput(0));

2.2 通過Parser載入已有網絡模型

  該方法也需要先創建IBuilder、INetworkDefinition，然後創建對應Caffe、UFF、ONNX的iParser，然後讀取解析模型。IBuilder是INetworkDefinition的工廠模式需要優先創建，不同的網絡框架對應的iParser輸出機制不同。

  對於Caffe：

// 定義IBuilder和iNetworkDefinition
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

// 創建Caffe的解析器
ICaffeParser* parser = createCaffeParser();

// 解析已有模型，DataType::kFLOAT替換爲DataType::kFLOAT能使用半精度
const IBlobNameToTensor* blobNameToTensor = parser->parse("deploy_file" , "modelFile", *network, DataType::kFLOAT);

// 指定網絡輸出
for (auto& s : outputs)
	network->markOutput(*blobNameToTensor->find(s.c_str()));

  對於TensorFlow模型使用UFF解析器（另有 TensorFlow-TensorRT 方法）：

// 定義IBuilder和iNetworkDefinition
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

// 創建UFF的解析器
IUFFParser* parser = createUffParser();

// 聲明網絡輸入與輸出，對TensorFlow輸入非CHW順序的需要先做轉換
parser->registerInput("Input_0", DimsCHW(1, 28, 28), UffInputOrder::kNCHW);
parser->registerOutput("Binary_3");

// 解析已有模型，DataType::kFLOAT替換爲DataType::kFLOAT能使用半精度
parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT);

  對於ONNX：

// 定義IBuilder和iNetworkDefinition
IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetwork();

// 創建ONNX的解析器
nvonnxparser::IONNXParser* parser = nvonnxparser::createONNXParser(*network, gLogger);

// 獲取模型
parser->parseFromFile(onnx_filename, ILogger::Severity::kWARNING);

三、創建引擎

  因爲引擎的優化是針對GPU硬件的，IBuilder需要放到同一個GPU上。IBuilder可以設置網絡運行的精度，也可以優化參數使性能最優，其主要的屬性是設置最大BatchSize和最大工作空間：

• maximum batch size，設定TensorRT對一批多大的數據進行優化，運行時可選擇較小的批大小；
• maximum workspace size，限定網絡中各層佔用存儲的最大空間，如果設置的太小，可能導致TensorRT無法實現某些層的優化。

// 創建引擎
builder->setMaxBatchSize(maxBatchSize);
builder->setMaxWorkspaceSize(1 << 20);
ICudaEngine* engine = builder->buildCudaEngine(*network);

// 完成處理後釋放
parser->destroy();
network->destroy();
builder->destroy();

四、序列化模型

  上一步建立的引擎可以直接使用，但考慮到創建引擎耗時過長，通常會選擇通過序列化創建二進制文件和反序列化讀取引擎，需要注意的是序列化後的文件不能跨平臺和TensorRT版本使用，因爲引擎是針對它們優化的。

// 通過引擎創建序列化文件
IHostMemory *serializedModel = engine->serialize();
// store model to disk
// <…>
serializedModel->destroy();

// 通過RunTime對象反序列化
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(modelData, modelSize, nullptr);

五、推理

// 創建context來開闢空間存儲中間值，一個engine可以有多個context來並行處理
IExecutionContext *context = engine->createExecutionContext();

// 指定輸入和輸出節點名來獲取輸入輸出索引
int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);

// 設置GPU上的緩存區陣列
void* buffers[2];
buffers[inputIndex] = inputbuffer;
buffers[outputIndex] = outputBuffer;

// 使用enqueue方法對CUDA內核在流上排隊，進行異步處理
context->enqueue(batchSize, buffers, stream, nullptr);
// 最後的參數是個可選CUDA事件，用於輸入緩存區被使用並可安全複用時發出信號

六、顯存管理

  TensorRT有兩種設備顯存管理機制：

• 默認方法，直接創建IExecutionContext會分配固定的顯存大小來保存激活數據，如果要避免直接分配，可以使用createExecutionContextWithoutDeviceMemory()，然後通過IExecutionContext::setDeviceMemory() 設置網絡運行需要的空間，通過ICudaEngine::getDeviceMemorySize()可以獲取顯存塊的大小；
• IGpuAllocator接口，使用IBuilder或IRuntime的setGpuAllocator(&allocator)來分配和釋放空間。

七、重置引擎

  TensorRT的引擎可以通過新的權重參數重置而不需要重新去創建，但引擎必須創建爲“refittable”，並且由於優化的原因更改一個參數可能需要連帶更新其它的參數。更新的參數可以通過refitter->getAll(...)查看

// 創建引擎前設置爲可重置類型的
...
builder->setRefittable(true);
builder->buildCudaEngine(network);

// 創建重置對象
ICudaEngine* engine = ...;
IRefitter* refitter = createInferRefitter(*engine,gLogger)

// 更新參數，以節點名爲“MyLayer”的卷積層爲例，新參數需要和原參數屬性相同
Weights newWeights = ...;
refitter->setWeights("MyLayer",WeightsRole::kKERNEL, newWeights);

// 查找關聯參數，先找到關聯參數數量，再查找它們的屬性
const int n = refitter->getMissing(0, nullptr, nullptr);
std::vector<const char*> layerNames(n);
std::vector<WeightsRole> weightsRoles(n);
refitter->getMissing(n, layerNames.data(), weightsRoles.data());

// 按屬性補充關聯參數，不補充額外的參數不會影響到更多的關聯參數
for (int i = 0; i < n; ++i)
	refitter->setWeights(layerNames[i], weightsRoles[i], Weights{...});

// 參數補充完成後，更新引擎，如果失敗則需查閱日誌看是否缺少參數未補充
bool success = refitter->refitCudaEngine();
assert(success);

// 釋放重置器
refitter->destroy();