TensorRT C++ Samples(1) sampleMNIST

---

0. 前言

• 目標：在根據官方文檔安裝完後嘗試測試一下TensorRT是否安裝成功。
• 代碼可以在 samples/sampleMNIST 中找到，也可以看 github 中對應路徑。
• TODO：進一步理解模型推理過程中的 stream/buffer/context 等變量的含義。

---

1. 運行過程

• 以下過程可以參考 README.md 相關信息。
• 下文中 ./ 指的是 tensorrt 所在路徑，如 /home/ubuntu/TensorRT-6.0.1.5。

1.1. 數據準備

• 第一步：在 MNIST官網 下載數據訓練數據 train-images-idx3-ubyte.gz 和 train-labels-idx1-ubyte.gz，並保存到 ./data/mnist 中。
• 第二步：通過 gzip -d xxx.gz 解壓上面兩個文件。
• 第三步：在 ./data/mnist/ 目錄下運行 python generate_pgms.py，生成若干 *.pgm 文件。

1.2. 代碼編譯與運行

• 第一步：在 ./samples/sampleMNIST 目錄下執行 make 命令。可執行文件生成在 ./bin/ 目錄下。
• 第二步：在 ./ 目錄下運行 ./bin/sample_mnist 即可看到預期結果。

&&&& RUNNING TensorRT.sample_mnist # ./bin/sample_mnist
[06/06/2020-16:26:42] [I] Building and running a GPU inference engine for MNIST
[06/06/2020-16:26:48] [I] [TRT] Detected 1 inputs and 1 output network tensors.
[06/06/2020-16:26:48] [I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@:.%@@@@@
@@@@@@@@@@@@@@@%%%%* .@@@@@@
@@@@@@@@@@@%-::.  =- :@@@@@@
@@@@@@@@@@@+    -.   *@@@@@@
@@@@@@@@@@=  -=@@#   @@@@@@@
@@@@@@@@@@= .%@@%-  +@@@@@@@
@@@@@@@@@@#. -@%:  +@@@@@@@@
@@@@@@@@@@@- .*.  +@@@@@@@@@
@@@@@@@@@@@=     -%@@@@@@@@@
@@@@@@@@@@@+    :%@@@@@@@@@@
@@@@@@@@@@@-   :@@@@@@@@@@@@
@@@@@@@@@@-    :@@@@@@@@@@@@
@@@@@@@@%+  :- :@@@@@@@@@@@@
@@@@@@@@*  +@* -@@@@@@@@@@@@
@@@@@@@#  *@%. =@@@@@@@@@@@@
@@@@@@@= :@@+  +@@@@@@@@@@@@
@@@@@@@= :@*  -@@@@@@@@@@@@@
@@@@@@@- -:  *@@@@@@@@@@@@@@
@@@@@@@+    =@@@@@@@@@@@@@@@
@@@@@@@@+ :+@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@

[06/06/2020-16:26:48] [I] Output:
0: 
1: 
2: 
3: 
4: 
5: 
6: 
7: 
8: **********
9: 

&&&& PASSED TensorRT.sample_mnist # ./bin/sample_mnist

---

2. 源碼解析

2.1. 基本概念

• 在幾乎TensorRT MNIST筆記的文章中都有一張圖：
  • 這張圖來自 TensorRT官方文檔，用於介紹TensorRT的基本流程，也就是下面源碼的基本流程。
    • 第一步：將訓練好的神經網絡模型轉換爲TensorRT的形式，並用TensorRT Optimizer進行優化。
    • 第二步：將在TensorRT Engine中運行優化好的TensorRT網絡結構。

2.2. 主函數

• 主要作用：
  • 解析輸入參數。
  • 構造 SampleMNIST 對象，調用相關方法實現Caffe模型轉換、TensorRT Engine推理。

int main(int argc, char** argv)
{
    // 使用 getopt_long 獲取參數
    // 如果所有參數都合法，則返回 true，否則返回 false
    samplesCommon::Args args;
    bool argsOK = samplesCommon::parseArgs(args, argc, argv);
    if (!argsOK)
    {
        gLogError << "Invalid arguments" << std::endl;
        printHelpInfo();
        return EXIT_FAILURE;
    }
    // 如果在輸入參數中指定了 help，那就返回
    if (args.help)
    {
        printHelpInfo();
        return EXIT_SUCCESS;
    }

    // 準備工作，如參數初始化
    auto sampleTest = gLogger.defineTest(gSampleName, argc, argv);
    gLogger.reportTestStart(sampleTest);
    samplesCommon::CaffeSampleParams params = initializeSampleParams(args);

    // 樣例主要就是通過 SampleMNIST 類來實現
    SampleMNIST sample(params);
    gLogInfo << "Building and running a GPU inference engine for MNIST" << std::endl;
    if (!sample.build())
    {
        return gLogger.reportFail(sampleTest);
    }
    if (!sample.infer())
    {
        return gLogger.reportFail(sampleTest);
    }
    
    // 結束工作流，釋放資源，返回結果
    if (!sample.teardown())
    {
        return gLogger.reportFail(sampleTest);
    }
    return gLogger.reportPass(sampleTest);
}

2.3. 將caffe模型轉換爲TensorRT可識別的形式

• 實現方式：通過主函數中的 sample.build() 實現。
• build函數基本流程
  • 主要工作就是，定義一個 network 對象，用於保存 caffe 模型轉換後的結果。
    • 通過 constructNetwork 實現。
  • 猜測所謂的 Tensor Optimizer 就是在 builder->buildEngineWithConfig 中實現的。

bool SampleMNIST::build()
{
    auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(gLogger.getTRTLogger()));
    if (!builder)
        return false;

    // 創建空的network，後面 constructNetwork 中會定義
    auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetwork());
    if (!network) return false;

    auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config) return false;

    // 創建caffe模型的parser，在constructNetwork函數中解析模型，轉換爲 network
    auto parser = SampleUniquePtr<nvcaffeparser1::ICaffeParser>(nvcaffeparser1::createCaffeParser());
    if (!parser) return false;

    // 解析 caffe 模型，並轉換爲 network 形式
    constructNetwork(parser, network);

    // 設置一些參數
    builder->setMaxBatchSize(mParams.batchSize);
    config->setMaxWorkspaceSize(16_MiB);
    config->setFlag(BuilderFlag::kGPU_FALLBACK);
    config->setFlag(BuilderFlag::kSTRICT_TYPES);
    if (mParams.fp16) config->setFlag(BuilderFlag::kFP16);
    if (mParams.int8) config->setFlag(BuilderFlag::kINT8);
    samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);

    // 構建 tensorrt 引擎
    // 注意，這個是成員變量
    mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config), samplesCommon::InferDeleter());
    if (!mEngine) return false;

    assert(network->getNbInputs() == 1);
    mInputDims = network->getInput(0)->getDimensions();
    assert(mInputDims.nbDims == 3);

    return true;
}

• caffe模型轉換具體實現過程
  • 其實模型轉換本身，parser->parse 一個函數就解決了。
  • 下面代碼的大量篇幅是在：在模型開頭添加 輸入圖片減去平均數 操作上。

void SampleMNIST::constructNetwork(SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser, SampleUniquePtr<nvinfer1::INetworkDefinition>& network)
{
    // 解析 caffe 的模型文件與權重文件，將結果寫入 network 中
    const nvcaffeparser1::IBlobNameToTensor* blobNameToTensor = parser->parse(
        mParams.prototxtFileName.c_str(),
        mParams.weightsFileName.c_str(),
        *network,
        nvinfer1::DataType::kFLOAT);

    // 標記模型輸出
    for (auto& s : mParams.outputTensorNames)
    {
        network->markOutput(*blobNameToTensor->find(s.c_str()));
    }

    // 在模型開頭添加 `輸入圖片減去平均數` 操作
    // add mean subtraction to the beginning of the network
    nvinfer1::Dims inputDims = network->getInput(0)->getDimensions();
    mMeanBlob = SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob>(parser->parseBinaryProto(mParams.meanFileName.c_str()));
    nvinfer1::Weights meanWeights{nvinfer1::DataType::kFLOAT, mMeanBlob->getData(), inputDims.d[1] * inputDims.d[2]};
    // For this sample, a large range based on the mean data is chosen and applied to the head of the network.
    // After the mean subtraction occurs, the range is expected to be between -127 and 127, so the rest of the network
    // is given a generic range.
    // The preferred method is use scales computed based on a representative data set
    // and apply each one individually based on the tensor. The range here is large enough for the
    // network, but is chosen for example purposes only.
    float maxMean = samplesCommon::getMaxValue(static_cast<const float*>(meanWeights.values), samplesCommon::volume(inputDims));
    // 模型中添加常量（圖片channel均值）
    auto mean = network->addConstant(nvinfer1::Dims3(1, inputDims.d[1], inputDims.d[2]), meanWeights);
    mean->getOutput(0)->setDynamicRange(-maxMean, maxMean);
    network->getInput(0)->setDynamicRange(-maxMean, maxMean);
    // 添加 減均值 操作
    auto meanSub = network->addElementWise(*network->getInput(0), *mean->getOutput(0), ElementWiseOperation::kSUB);
    meanSub->getOutput(0)->setDynamicRange(-maxMean, maxMean);
    network->getLayer(0)->setInput(0, *meanSub->getOutput(0));
    // 執行縮放，輸出結果爲 [-1, 1]
    samplesCommon::setAllTensorScales(network.get(), 127.0f, 127.0f);
}

2.4. 模型推理

• 主要工作：就是將轉換好的模型在tensorrt engine上跑一邊。
  • 裏面用到的各種東西比較多，目前也看不懂。
• 主要通過 infer() 函數完成，本函數的主要操作就是：
  • 讀取輸入數據(processInput)。
  • 通過 cuda stream / buffer 等進行推理。
  • 判斷輸出結果是否正確（verifyOutput）。

bool SampleMNIST::infer()
{
    // 實現具體推理過程

    // 緩存對象管理
    // Create RAII buffer manager object
    samplesCommon::BufferManager buffers(mEngine, mParams.batchSize);

    // 創建上下文
    auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
    if (!context)
    {
        return false;
    }

    // 隨機選擇一個數字
    // Pick a random digit to try to infer
    srand(time(NULL));
    const int digit = rand() % 10;

    // 讀取輸入數據到緩存對象中
    // 即將 digit 寫入 buffers 中，名字爲 mParams.inputTensorNames[0]
    // Read the input data into the managed buffers
    // There should be just 1 input tensor
    assert(mParams.inputTensorNames.size() == 1);
    if (!processInput(buffers, mParams.inputTensorNames[0], digit))
    {
        return false;
    }

    // 創建 cuda 流，準備執行推理
    // Create CUDA stream for the execution of this inference.
    cudaStream_t stream;
    CHECK(cudaStreamCreate(&stream));

    // 異步將數據從主機輸入緩衝區(buffer)複製到設備輸入緩衝區(stream)
    // Asynchronously copy data from host input buffers to device input buffers
    buffers.copyInputToDeviceAsync(stream);

    // 異步將推理任務加入隊列中
    // Asynchronously enqueue the inference work
    if (!context->enqueue(mParams.batchSize, buffers.getDeviceBindings().data(), stream, nullptr))
    {
        return false;
    }

    // 異步將模型結果從設備(stream)保存到主機緩衝區(buffers)
    // Asynchronously copy data from device output buffers to host output buffers
    buffers.copyOutputToHostAsync(stream);

    // 等待工作結束，關閉stream
    // Wait for the work in the stream to complete
    cudaStreamSynchronize(stream);
    // Release stream
    cudaStreamDestroy(stream);

    // 得到結果，判斷結果是否準確
    // 即從 buffer 中獲取名爲 mParams.outputTensorNames[0] 的結果，判斷與digit是否相同
    // Check and print the output of the inference
    // There should be just one output tensor
    assert(mParams.outputTensorNames.size() == 1);
    bool outputCorrect = verifyOutput(buffers, mParams.outputTensorNames[0], digit);

    return outputCorrect;
}

• 讀取輸入數據
  • 這部分沒啥好說的。

bool SampleMNIST::processInput(const samplesCommon::BufferManager& buffers, const std::string& inputTensorName, int inputFileIdx) const
{
    const int inputH = mInputDims.d[1];
    const int inputW = mInputDims.d[2];

    // Read a random digit file
    srand(unsigned(time(nullptr)));
    std::vector<uint8_t> fileData(inputH * inputW);
    readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);

    // Print ASCII representation of digit
    gLogInfo << "Input:\n";
    for (int i = 0; i < inputH * inputW; i++)
    {
        gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
    }
    gLogInfo << std::endl;

    float* hostInputBuffer = static_cast<float*>(buffers.getHostBuffer(inputTensorName));

    for (int i = 0; i < inputH * inputW; i++)
    {
        hostInputBuffer[i] = float(fileData[i]);
    }

    return true;
}

• 驗證輸出結果
  • 還會輸出可視化結果，有種夢迴當年的感覺。

bool SampleMNIST::verifyOutput(const samplesCommon::BufferManager& buffers, const std::string& outputTensorName, int groundTruthDigit) const
{
    // 獲取 host buffer 中的輸出tensor數值
    // 應該是10個數字的概率
    const float* prob = static_cast<const float*>(buffers.getHostBuffer(outputTensorName));

    // Print histogram of the output distribution
    gLogInfo << "Output:\n";
    float val{0.0f};
    int idx{0};
    const int kDIGITS = 10;

    for (int i = 0; i < kDIGITS; i++)
    {
        if (val < prob[i])
        {
            val = prob[i];
            idx = i;
        }

        gLogInfo << i << ": " << std::string(int(std::floor(prob[i] * 10 + 0.5f)), '*') << "\n";
    }
    gLogInfo << std::endl;

    return (idx == groundTruthDigit && val > 0.9f);
}