devbox
程序诗人
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

TensorRT学习笔记--官方示例sampleMNIST.cpp的理解与运行_sample_mnist

作者：程序诗人 | 2024-02-03 18:21:45

赞

踩

sample_mnist

目录

1--前言

2--编写CMakeLists.txt

3--代码注释

4--运行结果

1--前言

        基于 TensorRT 8.2.5.1 和 Cuda11.3，在 Ubuntu 20.04 编写 CMakeLists.txt 进行编译，生成可执行文件；

2--编写CMakeLists.txt

        提供一个 CMakeLists.txt 编写样例，需根据个人实际修改具体路径：

  1. cmake_minimum_required(VERSION 3.13)
  2. project(TensorRT_test)
  3. set(CMAKE_CXX_STANDARD 11)
  4.  
  5. set(SAMPLES_COMMON_SOURCES "/home/liujinfu/Downloads/TensorRT-8.2.5.1/samples/common/logger.cpp")
  6. add_executable(TensorRT_test sampleMNIST.cpp ${SAMPLES_COMMON_SOURCES})
  7.  
  8. # add TensorRT8
  9. include_directories(/home/liujinfu/Downloads/TensorRT-8.2.5.1/include)
  10. include_directories(/home/liujinfu/Downloads/TensorRT-8.2.5.1/samples/common)
  11. set(TENSORRT_LIB_PATH "/home/liujinfu/Downloads/TensorRT-8.2.5.1/lib")
  12. file(GLOB LIBS "${TENSORRT_LIB_PATH}/*.so")
  13.  
  14. # add CUDA
  15. find_package(CUDA 11.3 REQUIRED)
  16. message("CUDA_LIBRARIES:${CUDA_LIBRARIES}")
  17. message("CUDA_INCLUDE_DIRS:${CUDA_INCLUDE_DIRS}")
  18. include_directories(${CUDA_INCLUDE_DIRS})
  19.  
  20. # link
  21. target_link_libraries(TensorRT_test ${LIBS} ${CUDA_LIBRARIES})

3--代码注释

  1. // 包含头文件
  2. #include "argsParser.h"
  3. #include "buffers.h"
  4. #include "common.h"
  5. #include "logger.h"
  6.  
  7. #include "NvCaffeParser.h"
  8. #include "NvInfer.h"
  9.  
  10. #include <algorithm>
  11. #include <cmath>
  12. #include <cuda_runtime_api.h>
  13. #include <fstream>
  14. #include <iostream>
  15. #include <sstream>
  16.  
  17. using samplesCommon::SampleUniquePtr;
  18.  
  19. const std::string gSampleName = "TensorRT.sample_mnist";
  20.  
  21.  
  22. // 定义 SampleMNIST 类
  23. class SampleMNIST
  24. {
  25. public:
  26.     // 构造函数
  27.     SampleMNIST(const samplesCommon::CaffeSampleParams& params)
  28.         : mParams(params)
  29.     {
  30.     }
  31.  
  32.     // 用于 build engine 的成员函数
  33.     bool build();
  34.  
  35.     // 用于 inference 的成员函数
  36.     bool infer();
  37.  
  38.     // 用于 清理状态 的成员函数
  39.     bool teardown();
  40.  
  41. private:
  42.     // 使用 caffe 的 parser 创建网络并标记输出层
  43.     bool constructNetwork(
  44.         SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser, SampleUniquePtr<nvinfer1::INetworkDefinition>& network);
  45.  
  46.     // 对输入数据进行前处理
  47.     bool processInput(
  48.         const samplesCommon::BufferManager& buffers, const std::string& inputTensorName, int inputFileIdx) const;
  49.  
  50.     // 验证输出结果是否正确并打印
  51.     bool verifyOutput(
  52.         const samplesCommon::BufferManager& buffers, const std::string& outputTensorName, int groundTruthDigit) const;
  53.  
  54.     std::shared_ptr<nvinfer1::ICudaEngine> mEngine{nullptr}; //!< The TensorRT engine used to run the network
  55.  
  56.     samplesCommon::CaffeSampleParams mParams; //!< The parameters for the sample.
  57.  
  58.     nvinfer1::Dims mInputDims; // 输入数据的维度
  59.  
  60.     SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob>
  61.         mMeanBlob; //! the mean blob, which we need to keep around until build is done
  62. };
  63.  
  64.  
  65. // build engine 的成员函数实现
  66. bool SampleMNIST::build()
  67. {
  68.     // 创建 builder
  69.     auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
  70.     if (!builder)
  71.     {
  72.         return false;
  73.     }
  74.  
  75.     // 使用 builder 创建 network
  76.     auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(0));
  77.     if (!network)
  78.     {
  79.         return false;
  80.     }
  81.  
  82.     // 创建config，用于配置模型
  83.     auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
  84.     if (!config)
  85.     {
  86.         return false;
  87.     }
  88.  
  89.     // 创建parser，用于解析模型
  90.     auto parser = SampleUniquePtr<nvcaffeparser1::ICaffeParser>(nvcaffeparser1::createCaffeParser());
  91.     if (!parser)
  92.     {
  93.         return false;
  94.     }
  95.  
  96.     // 执行 constructNetwork() 成员函数，使用 parser 创建 MNIST network 并标记输出层
  97.     if (!constructNetwork(parser, network))
  98.     {
  99.         return false;
  100.     }
  101.  
  102.     // 设置推理网络的相关参数
  103.     builder->setMaxBatchSize(mParams.batchSize); // 设定 batchsize
  104.     config->setMaxWorkspaceSize(16_MiB); // 设定 workspace
  105.     config->setFlag(BuilderFlag::kGPU_FALLBACK); // 启用GPU回退模式
  106.     // 设置精度计算
  107.     if (mParams.fp16) // 使用fp16精度
  108.     {
  109.         config->setFlag(BuilderFlag::kFP16);
  110.     }
  111.     if (mParams.int8) // 使用int8精度
  112.     {
  113.         config->setFlag(BuilderFlag::kINT8);
  114.     }
  115.  
  116.     // 启用 DLA
  117.     samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);
  118.  
  119.     // 创建 Cuda stream
  120.     auto profileStream = samplesCommon::makeCudaStream();
  121.     if (!profileStream)
  122.     {
  123.         return false;
  124.     }
  125.     config->setProfileStream(*profileStream);
  126.  
  127.     // builder->buildSerializedNetwork 创建推理引擎
  128.     SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
  129.     if (!plan)
  130.     {
  131.         return false;
  132.     }
  133.  
  134.     // 序列化模型后，创建 Runtime 接口
  135.     SampleUniquePtr<IRuntime> runtime{createInferRuntime(sample::gLogger.getTRTLogger())};
  136.     if (!runtime)
  137.     {
  138.         return false;
  139.     }
  140.  
  141.     // 反序列化
  142.     mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
  143.         runtime->deserializeCudaEngine(plan->data(), plan->size()), samplesCommon::InferDeleter());
  144.     if (!mEngine)
  145.     {
  146.         return false;
  147.     }
  148.  
  149.     // 断定输入的 batchsize 为1，即单次只能处理一个样本
  150.     ASSERT(network->getNbInputs() == 1);
  151.     // 获取输入的维度，断定输入数据的维度为3
  152.     mInputDims = network->getInput(0)->getDimensions();
  153.     ASSERT(mInputDims.nbDims == 3);
  154.  
  155.     // build engine 成功，返回True
  156.     return true;
  157. }
  158.  
  159. // 输入数据前处理成员函数的实现
  160. bool SampleMNIST::processInput(
  161.     const samplesCommon::BufferManager& buffers, const std::string& inputTensorName, int inputFileIdx) const
  162. {
  163.     // 获取输入图像的高和宽，mInputDims 已在私有成员变量中定义
  164.     const int inputH = mInputDims.d[1];
  165.     const int inputW = mInputDims.d[2];
  166.  
  167.     // 随机读取一个数字文件
  168.     srand(unsigned(time(nullptr)));
  169.     std::vector<uint8_t> fileData(inputH * inputW);
  170.     readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);
  171.  
  172.     // 打印数字的 ASCII 表示，在控制台中用 ASCII 码显示图片
  173.     sample::gLogInfo << "Input:\n";
  174.     for (int i = 0; i < inputH * inputW; i++)
  175.     {
  176.         sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
  177.     }
  178.     sample::gLogInfo << std::endl;
  179.  
  180.     // 将输入数据存储到主机(host)的内存中
  181.     float* hostInputBuffer = static_cast<float*>(buffers.getHostBuffer(inputTensorName));
  182.  
  183.     for (int i = 0; i < inputH * inputW; i++)
  184.     {
  185.         hostInputBuffer[i] = float(fileData[i]);
  186.     }
  187.  
  188.     return true;
  189. }
  190.  
  191. // 成员函数实现，验证输出结果是否正确并打印
  192. bool SampleMNIST::verifyOutput(
  193.     const samplesCommon::BufferManager& buffers, const std::string& outputTensorName, int groundTruthDigit) const
  194. {
  195.     // 从 host 的 output buffer 中读取 输出结果
  196.     const float* prob = static_cast<const float*>(buffers.getHostBuffer(outputTensorName));
  197.  
  198.     // 打印输出分布的直方图
  199.     sample::gLogInfo << "Output:\n";
  200.     float val{0.0f};
  201.     int idx{0};
  202.     const int kDIGITS = 10;
  203.  
  204.     for (int i = 0; i < kDIGITS; i++)
  205.     {
  206.         if (val < prob[i])
  207.         {
  208.             val = prob[i];
  209.             idx = i;
  210.         }
  211.  
  212.         sample::gLogInfo << i << ": " << std::string(int(std::floor(prob[i] * 10 + 0.5f)), '*') << "\n";
  213.     }
  214.     sample::gLogInfo << std::endl;
  215.  
  216.     return (idx == groundTruthDigit && val > 0.9f);
  217. }
  218.  
  219. // constructNetwork()成员函数的实现
  220. bool SampleMNIST::constructNetwork(
  221.     SampleUniquePtr<nvcaffeparser1::ICaffeParser>& parser, SampleUniquePtr<nvinfer1::INetworkDefinition>& network)
  222. {
  223.     const nvcaffeparser1::IBlobNameToTensor* blobNameToTensor = parser->parse(
  224.         mParams.prototxtFileName.c_str(), mParams.weightsFileName.c_str(), *network, nvinfer1::DataType::kFLOAT);
  225.  
  226.     // 输出 Tensor 标记
  227.     for (auto& s : mParams.outputTensorNames)
  228.     {
  229.         network->markOutput(*blobNameToTensor->find(s.c_str()));
  230.     }
  231.  
  232.     // 在网络开头添加减均值的操作（针对本示例而言）
  233.     nvinfer1::Dims inputDims = network->getInput(0)->getDimensions();
  234.     mMeanBlob
  235.         = SampleUniquePtr<nvcaffeparser1::IBinaryProtoBlob>(parser->parseBinaryProto(mParams.meanFileName.c_str()));
  236.     nvinfer1::Weights meanWeights{nvinfer1::DataType::kFLOAT, mMeanBlob->getData(), inputDims.d[1] * inputDims.d[2]};
  237.  
  238.     float maxMean
  239.         = samplesCommon::getMaxValue(static_cast<const float*>(meanWeights.values), samplesCommon::volume(inputDims));
  240.  
  241.     auto mean = network->addConstant(nvinfer1::Dims3(1, inputDims.d[1], inputDims.d[2]), meanWeights);
  242.     if (!mean->getOutput(0)->setDynamicRange(-maxMean, maxMean))
  243.     {
  244.         return false;
  245.     }
  246.     if (!network->getInput(0)->setDynamicRange(-maxMean, maxMean))
  247.     {
  248.         return false;
  249.     }
  250.     auto meanSub = network->addElementWise(*network->getInput(0), *mean->getOutput(0), ElementWiseOperation::kSUB);
  251.     if (!meanSub->getOutput(0)->setDynamicRange(-maxMean, maxMean))
  252.     {
  253.         return false;
  254.     }
  255.     network->getLayer(0)->setInput(0, *meanSub->getOutput(0));
  256.     samplesCommon::setAllDynamicRanges(network.get(), 127.0f, 127.0f);
  257.  
  258.     return true;
  259. }
  260.  
  261. // 推理过程的成员函数实现
  262. bool SampleMNIST::infer()
  263. {
  264.     // 创建 buffer 管理对象
  265.     samplesCommon::BufferManager buffers(mEngine, mParams.batchSize);
  266.  
  267.     // 创建上下文 Context
  268.     auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
  269.     if (!context)
  270.     {
  271.         return false;
  272.     }
  273.  
  274.     // 随机选择一个数字进行推理
  275.     srand(time(NULL));
  276.     const int digit = rand() % 10;
  277.  
  278.     // 断定数字的个数为1
  279.     ASSERT(mParams.inputTensorNames.size() == 1);
  280.     // 对输入数字进行前处理
  281.     if (!processInput(buffers, mParams.inputTensorNames[0], digit))
  282.     {
  283.         return false;
  284.     }
  285.     
  286.     // 创建 Cuda stream 
  287.     cudaStream_t stream;
  288.     CHECK(cudaStreamCreate(&stream));
  289.  
  290.     // 异步地将 input_data 从 host 的 input buffer 拷贝到 device 的 input buffer （host -> device）
  291.     buffers.copyInputToDeviceAsync(stream);
  292.  
  293.     // 异步地执行推理
  294.     if (!context->enqueue(mParams.batchSize, buffers.getDeviceBindings().data(), stream, nullptr))
  295.     {
  296.         return false;
  297.     }
  298.     
  299.     // 推理结果从 device 转移到 host
  300.     buffers.copyOutputToHostAsync(stream);
  301.  
  302.     // 等待 stream 中的 work 完成
  303.     cudaStreamSynchronize(stream);
  304.  
  305.     // 释放stream
  306.     cudaStreamDestroy(stream);
  307.  
  308.     // 断定输出结果只有一个Tensor
  309.     ASSERT(mParams.outputTensorNames.size() == 1);
  310.     // 执行函数 verifyOutput() 验证结果的正确性，并打印推理结果，这个函数可以理解成后处理函数
  311.     bool outputCorrect = verifyOutput(buffers, mParams.outputTensorNames[0], digit);
  312.  
  313.     return outputCorrect;
  314. }
  315.  
  316. // 成员函数实现，清理对象的所有状态
  317. bool SampleMNIST::teardown()
  318. {
  319.     nvcaffeparser1::shutdownProtobufLibrary();
  320.     return true;
  321. }
  322.  
  323. // 使用命令行参数初始化结构体参数成员
  324. samplesCommon::CaffeSampleParams initializeSampleParams(const samplesCommon::Args& args)
  325. {
  326.     samplesCommon::CaffeSampleParams params;
  327.     if (args.dataDirs.empty()) // 输入参数的路径为空时，使用默认的路径
  328.     {
  329.         params.dataDirs.push_back("data/mnist/");
  330.         params.dataDirs.push_back("data/samples/mnist/");
  331.     }
  332.     else // 输入参数的路径不为空时，使用输入的路径
  333.     {
  334.         params.dataDirs = args.dataDirs;
  335.     }
  336.  
  337.     // 从路径中读取相应的文件
  338.     params.prototxtFileName = locateFile("mnist.prototxt", params.dataDirs);
  339.     params.weightsFileName = locateFile("mnist.caffemodel", params.dataDirs);
  340.     params.meanFileName = locateFile("mnist_mean.binaryproto", params.dataDirs);
  341.     params.inputTensorNames.push_back("data");
  342.     params.batchSize = 1;
  343.     params.outputTensorNames.push_back("prob");
  344.     params.dlaCore = args.useDLACore;
  345.     params.int8 = args.runInInt8;
  346.     params.fp16 = args.runInFp16;
  347.  
  348.     return params;
  349. }
  350.  
  351.  
  352. // 定义打印参数含义的函数
  353. void printHelpInfo()
  354. {
  355.     std::cout
  356.         << "Usage: ./sample_mnist [-h or --help] [-d or --datadir=<path to data directory>] [--useDLACore=<int>]\n";
  357.     std::cout << "--help          Display help information\n";
  358.     std::cout << "--datadir       Specify path to a data directory, overriding the default. This option can be used "
  359.                  "multiple times to add multiple directories. If no data directories are given, the default is to use "
  360.                  "(data/samples/mnist/, data/mnist/)"
  361.               << std::endl;
  362.     std::cout << "--useDLACore=N  Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, "
  363.                  "where n is the number of DLA engines on the platform."
  364.               << std::endl;
  365.     std::cout << "--int8          Run in Int8 mode.\n";
  366.     std::cout << "--fp16          Run in FP16 mode.\n";
  367. }
  368.  
  369. int main(int argc, char** argv)
  370. {   
  371.     // 解析参数的含义，例如--help
  372.     samplesCommon::Args args;
  373.     bool argsOK = samplesCommon::parseArgs(args, argc, argv);
  374.  
  375.     // 参数不匹配
  376.     if (!argsOK)
  377.     {
  378.         sample::gLogError << "Invalid arguments" << std::endl;
  379.         printHelpInfo();
  380.         return EXIT_FAILURE;
  381.     }
  382.  
  383.     // 输出参数帮助信息
  384.     if (args.help)
  385.     {
  386.         printHelpInfo();
  387.         return EXIT_SUCCESS;
  388.     }
  389.  
  390.     // 定义用于 test 的logger
  391.     auto sampleTest = sample::gLogger.defineTest(gSampleName, argc, argv);
  392.  
  393.     // 输出: 表明 test 开始
  394.     sample::gLogger.reportTestStart(sampleTest);
  395.  
  396.     // 执行定义的 initializeSampleParams() 函数，使用命令行参数初始化结构体参数成员
  397.     samplesCommon::CaffeSampleParams params = initializeSampleParams(args);
  398.  
  399.     // 使用 SampleMNIST 类根据params定义一个sample对象
  400.     SampleMNIST sample(params);
  401.     sample::gLogInfo << "Building and running a GPU inference engine for MNIST" << std::endl; // 打印信息
  402.  
  403.     // 执行 sample.build() 编译推理引擎
  404.     if (!sample.build())
  405.     {
  406.         return sample::gLogger.reportFail(sampleTest); // 编译失败，使用gLogger报告信息
  407.     }
  408.  
  409.     // 执行 sample.infer() 进行 inference 推理
  410.     if (!sample.infer())
  411.     {
  412.         return sample::gLogger.reportFail(sampleTest); // 推理失败，使用gLogger报告信息
  413.     }
  414.  
  415.     // 清除 sample 的所有状态，释放内存
  416.     if (!sample.teardown())
  417.     {
  418.         return sample::gLogger.reportFail(sampleTest);
  419.     }
  420.  
  421.     return sample::gLogger.reportPass(sampleTest);
  422. }

4--运行结果

        执行以下命令进行编译，生成可执行文件：

  1. mkdir build
  2.  
  3. cd build
  4.  
  5. cmake ..
  6.  
  7. make

        执行以下命令运行可执行文件，并指定模型地址：

./TensorRT_test -d /home/liujinfu/Downloads/TensorRT-8.2.5.1/data/mnist

        运行结果截图：

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/article/detail/57543
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号