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

TensorRT学习笔记--官方示例sampleOnnxMNIST.cpp的理解与运行_params.inputtensornames.push_back

作者：编程变革者 | 2024-02-03 18:26:19

赞

踩

params.inputtensornames.push_back

前言

        重点关注官方示例如何导入Onnx模型，相关API的用法；

1--代码

  1. #include "argsParser.h"
  2. #include "buffers.h"
  3. #include "common.h"
  4. #include "logger.h"
  5. #include "parserOnnxConfig.h"
  6.  
  7. #include "NvInfer.h"
  8. #include <cuda_runtime_api.h>
  9.  
  10. #include <cstdlib>
  11. #include <fstream>
  12. #include <iostream>
  13. #include <sstream>
  14.  
  15. using samplesCommon::SampleUniquePtr;
  16.  
  17. const std::string gSampleName = "TensorRT.sample_onnx_mnist";
  18.  
  19. class SampleOnnxMNIST
  20. {
  21. public:
  22.     SampleOnnxMNIST(const samplesCommon::OnnxSampleParams& params): mParams(params), mEngine(nullptr){}
  23.  
  24.     bool build();
  25.  
  26.     bool infer();
  27.  
  28. private:
  29.     samplesCommon::OnnxSampleParams mParams; 
  30.  
  31.     nvinfer1::Dims mInputDims;  // 网络的输入维度
  32.     nvinfer1::Dims mOutputDims; // 网络的输出维度
  33.     int mNumber{0};             // 类别数
  34.  
  35.     std::shared_ptr<nvinfer1::ICudaEngine> mEngine; // 初始化engine
  36.  
  37.     bool constructNetwork(SampleUniquePtr<nvinfer1::IBuilder>& builder,
  38.         SampleUniquePtr<nvinfer1::INetworkDefinition>& network, SampleUniquePtr<nvinfer1::IBuilderConfig>& config,
  39.         SampleUniquePtr<nvonnxparser::IParser>& parser);
  40.  
  41.     // 前处理
  42.     bool processInput(const samplesCommon::BufferManager& buffers);
  43.  
  44.     // 验证结果
  45.     bool verifyOutput(const samplesCommon::BufferManager& buffers);
  46. };
  47.  
  48. bool SampleOnnxMNIST::build()
  49. {
  50.     // 创建 builder
  51.     auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
  52.     if (!builder)
  53.     {
  54.         return false;
  55.     }
  56.  
  57.     // 显式 batch
  58.     const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
  59.     // 创建network
  60.     auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
  61.     if (!network)
  62.     {
  63.         return false;
  64.     }
  65.     // 创建config
  66.     auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
  67.     if (!config)
  68.     {
  69.         return false;
  70.     }
  71.     // 创建parser
  72.     auto parser = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
  73.     if (!parser)
  74.     {
  75.         return false;
  76.     }
  77.  
  78.     // 调用成员函数生成网络
  79.     auto constructed = constructNetwork(builder, network, config, parser);
  80.     if (!constructed)
  81.     {
  82.         return false;
  83.     }
  84.  
  85.     // CUDA stream used for profiling by the builder.
  86.     auto profileStream = samplesCommon::makeCudaStream(); // 创建Cuda stream
  87.     if (!profileStream)
  88.     {
  89.         return false;
  90.     }
  91.     config->setProfileStream(*profileStream);
  92.  
  93.     SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)}; // 创建推理引擎
  94.     if (!plan)
  95.     {
  96.         return false;
  97.     }
  98.  
  99.     SampleUniquePtr<IRuntime> runtime{createInferRuntime(sample::gLogger.getTRTLogger())}; // 创建Runtime接口
  100.     if (!runtime)
  101.     {
  102.         return false;
  103.     }
  104.  
  105.     mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
  106.         runtime->deserializeCudaEngine(plan->data(), plan->size()), samplesCommon::InferDeleter());
  107.     if (!mEngine)
  108.     {
  109.         return false;
  110.     }
  111.  
  112.     ASSERT(network->getNbInputs() == 1); // 输入 batch 为1
  113.     mInputDims = network->getInput(0)->getDimensions();
  114.     ASSERT(mInputDims.nbDims == 4); // 输入维度为 4（包含 batch 维度）
  115.  
  116.     ASSERT(network->getNbOutputs() == 1); // 输出 batch 为1
  117.     mOutputDims = network->getOutput(0)->getDimensions();
  118.     ASSERT(mOutputDims.nbDims == 2); // 输出维度为 2
  119.  
  120.     return true;
  121. }
  122.  
  123. // constructNetwork 成员函数实现
  124. bool SampleOnnxMNIST::constructNetwork(SampleUniquePtr<nvinfer1::IBuilder>& builder,
  125.     SampleUniquePtr<nvinfer1::INetworkDefinition>& network, SampleUniquePtr<nvinfer1::IBuilderConfig>& config,
  126.     SampleUniquePtr<nvonnxparser::IParser>& parser)
  127. {
  128.     auto parsed = parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
  129.         static_cast<int>(sample::gLogger.getReportableSeverity())); // 解析onnx模型
  130.     if (!parsed)
  131.     {
  132.         return false;
  133.     }
  134.  
  135.     config->setMaxWorkspaceSize(16_MiB); // 设置最大工作空间
  136.     if (mParams.fp16) // 设置精度
  137.     {
  138.         config->setFlag(BuilderFlag::kFP16);
  139.     }
  140.     if (mParams.int8)
  141.     {
  142.         config->setFlag(BuilderFlag::kINT8);
  143.         samplesCommon::setAllDynamicRanges(network.get(), 127.0f, 127.0f);
  144.     }
  145.  
  146.     samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);
  147.  
  148.     return true;
  149. }
  150.  
  151. bool SampleOnnxMNIST::infer()
  152. {
  153.     // 创建内存管理
  154.     samplesCommon::BufferManager buffers(mEngine);
  155.  
  156.     // 创建context
  157.     auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
  158.     if (!context)
  159.     {
  160.         return false;
  161.     }
  162.  
  163.     // 从 buffer 中读取数据，并进行前处理
  164.     ASSERT(mParams.inputTensorNames.size() == 1);
  165.     if (!processInput(buffers))
  166.     {
  167.         return false;
  168.     }
  169.  
  170.     // 将数据复制到 GPU 中
  171.     buffers.copyInputToDevice();
  172.  
  173.     bool status = context->executeV2(buffers.getDeviceBindings().data()); // 执行推理
  174.     if (!status)
  175.     {
  176.         return false;
  177.     }
  178.  
  179.     // 将结果复制到 CPU 中
  180.     buffers.copyOutputToHost();
  181.  
  182.     // 验证推理结果
  183.     if (!verifyOutput(buffers))
  184.     {
  185.         return false;
  186.     }
  187.  
  188.     return true;
  189. }
  190.  
  191. // 前处理成员函数实现
  192. bool SampleOnnxMNIST::processInput(const samplesCommon::BufferManager& buffers) // 传入 buffers 的引用
  193. {
  194.     const int inputH = mInputDims.d[2];
  195.     const int inputW = mInputDims.d[3];
  196.  
  197.     // 随机选择一个数据进行读取
  198.     srand(unsigned(time(nullptr)));
  199.     std::vector<uint8_t> fileData(inputH * inputW);
  200.     mNumber = rand() % 10;
  201.     readPGMFile(locateFile(std::to_string(mNumber) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);
  202.  
  203.     // Print an ascii representation
  204.     sample::gLogInfo << "Input:" << std::endl;
  205.     for (int i = 0; i < inputH * inputW; i++)
  206.     {
  207.         sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
  208.     }
  209.     sample::gLogInfo << std::endl;
  210.  
  211.     // buffers.getHostBuffer() 返回 Name 对应的 buffer 地址
  212.     float* hostDataBuffer = static_cast<float*>(buffers.getHostBuffer(mParams.inputTensorNames[0])); 
  213.     for (int i = 0; i < inputH * inputW; i++) // 对这段 buffer 地址的内容进行赋值操作，由于 processInput() 函数传入的是buffer引用，所以能改变buffer的值
  214.     {
  215.         hostDataBuffer[i] = 1.0 - float(fileData[i] / 255.0);
  216.     }
  217.  
  218.     return true;
  219. }
  220.  
  221. // 验证推理结果的成员函数实现
  222. bool SampleOnnxMNIST::verifyOutput(const samplesCommon::BufferManager& buffers)
  223. {
  224.     const int outputSize = mOutputDims.d[1];
  225.     float* output = static_cast<float*>(buffers.getHostBuffer(mParams.outputTensorNames[0]));
  226.     float val{0.0f};
  227.     int idx{0};
  228.  
  229.     // Calculate Softmax
  230.     float sum{0.0f};
  231.     for (int i = 0; i < outputSize; i++)
  232.     {
  233.         output[i] = exp(output[i]);
  234.         sum += output[i];
  235.     }
  236.  
  237.     sample::gLogInfo << "Output:" << std::endl;
  238.     for (int i = 0; i < outputSize; i++)
  239.     {
  240.         output[i] /= sum;
  241.         val = std::max(val, output[i]);
  242.         if (val == output[i])
  243.         {
  244.             idx = i;
  245.         }
  246.  
  247.         sample::gLogInfo << " Prob " << i << "  " << std::fixed << std::setw(5) << std::setprecision(4) << output[i]
  248.                          << " "
  249.                          << "Class " << i << ": " << std::string(int(std::floor(output[i] * 10 + 0.5f)), '*')
  250.                          << std::endl;
  251.     }
  252.     sample::gLogInfo << std::endl;
  253.  
  254.     return idx == mNumber && val > 0.9f;
  255. }
  256.  
  257. // 初始化参数
  258. samplesCommon::OnnxSampleParams initializeSampleParams(const samplesCommon::Args& args)
  259. {
  260.     samplesCommon::OnnxSampleParams params; // 创建参数对象
  261.     if (args.dataDirs.empty()) //!< Use default directories if user hasn't provided directory paths
  262.     {
  263.         params.dataDirs.push_back("data/mnist/");
  264.         params.dataDirs.push_back("data/samples/mnist/");
  265.     }
  266.     else //!< Use the data directory provided by the user
  267.     {
  268.         params.dataDirs = args.dataDirs;
  269.     }
  270.     // 设置参数对象的默认属性
  271.     params.onnxFileName = "mnist.onnx";
  272.     params.inputTensorNames.push_back("Input3");
  273.     params.outputTensorNames.push_back("Plus214_Output_0");
  274.     params.dlaCore = args.useDLACore;
  275.     params.int8 = args.runInInt8;
  276.     params.fp16 = args.runInFp16;
  277.  
  278.     return params;
  279. }
  280.  
  281. // 打印参数信息
  282. void printHelpInfo()
  283. {
  284.     std::cout
  285.         << "Usage: ./sample_onnx_mnist [-h or --help] [-d or --datadir=<path to data directory>] [--useDLACore=<int>]"
  286.         << std::endl;
  287.     std::cout << "--help          Display help information" << std::endl;
  288.     std::cout << "--datadir       Specify path to a data directory, overriding the default. This option can be used "
  289.                  "multiple times to add multiple directories. If no data directories are given, the default is to use "
  290.                  "(data/samples/mnist/, data/mnist/)"
  291.               << std::endl;
  292.     std::cout << "--useDLACore=N  Specify a DLA engine for layers that support DLA. Value can range from 0 to n-1, "
  293.                  "where n is the number of DLA engines on the platform."
  294.               << std::endl;
  295.     std::cout << "--int8          Run in Int8 mode." << std::endl;
  296.     std::cout << "--fp16          Run in FP16 mode." << std::endl;
  297. }
  298.  
  299. int main(int argc, char** argv)
  300. {
  301.     samplesCommon::Args args;
  302.     bool argsOK = samplesCommon::parseArgs(args, argc, argv);
  303.     if (!argsOK)
  304.     {
  305.         sample::gLogError << "Invalid arguments" << std::endl;
  306.         printHelpInfo();
  307.         return EXIT_FAILURE;
  308.     }
  309.     if (args.help)
  310.     {
  311.         printHelpInfo();
  312.         return EXIT_SUCCESS;
  313.     }
  314.  
  315.     auto sampleTest = sample::gLogger.defineTest(gSampleName, argc, argv);
  316.  
  317.     sample::gLogger.reportTestStart(sampleTest);
  318.  
  319.     SampleOnnxMNIST sample(initializeSampleParams(args)); // 创建对象
  320.  
  321.     sample::gLogInfo << "Building and running a GPU inference engine for Onnx MNIST" << std::endl;
  322.  
  323.     if (!sample.build())
  324.     {
  325.         return sample::gLogger.reportFail(sampleTest);
  326.     }
  327.     if (!sample.infer())
  328.     {
  329.         return sample::gLogger.reportFail(sampleTest);
  330.     }
  331.  
  332.     return sample::gLogger.reportPass(sampleTest);
  333. }

2--编译

① 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_OnnxMNIST sampleOnnxMNIST.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_OnnxMNIST ${LIBS} ${CUDA_LIBRARIES})

② 编译

  1. mkdir build && cd build
  2.  
  3. cmake ..
  4.  
  5. make

3--运行结果

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

 

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop】
推荐阅读
相关标签

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

  

闽ICP备14008679号