TensorRT 入门Hello World实战解析sampleOnnxMNIST 从 ONNX 模型到推理引擎的完整流程【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT本篇技术指南以 samples/sampleOnnxMNIST/README.md 为骨架结合仓库中的 sampleOnnxMNIST.cpp 源码完整剖析 NVIDIA TensorRT 官方Hello World示例 sampleOnnxMNIST它如何把一份训练自 MNIST 数据集的 ONNX 模型解析为 TensorRT 网络、构建推理引擎并执行分类推理。读完本文你将掌握 ONNX Parser 的初始化与调用方式、基于IBuilder/IBuilderConfig的引擎构建流程、基于IExecutionContext的推理执行链路以及该示例全部命令行参数的用法为后续上手 sampleDynamicReshape、sampleIOFormats 等更复杂的 ONNX 示例打下基础。一、示例概览为什么它是 TensorRT 的Hello WorldsampleOnnxMNIST 是 TensorRT 仓库中最经典的入门示例其地位可以从 samples/README.md 的目录组织看出——它被归类在Hello World Samples 的第一位描述为 Hello World For TensorRT With ONNX与 Python 侧的 network_api_pytorch_mnist基于INetworkAPI形成 C/Python 双入口的对应关系。该示例的核心任务是将一个基于 MNIST 数据集训练、以Open Neural Network ExchangeONNX格式保存的模型转换为 TensorRT 网络并在该网络上运行推理。ONNX 是一种用于表示深度学习模型的标准格式它让模型可以在不同深度学习框架之间迁移。对 TensorRT 用户而言ONNX 是连接训练框架PyTorch / TensorFlow / CNTK 等与TensorRT 推理引擎之间最常见的桥梁之一而 sampleOnnxMNIST 正是展示这条桥怎么走的最小闭环。具体来说示例完成了三件事将 ONNX 模型转换为 TensorRT 网络通过 ONNX Parser 读取.onnx文件并填充INetworkDefinition构建引擎利用IBuilder与IBuilderConfig将网络编译为序列化的推理引擎serialized engine运行推理使用IExecutionContext在 GPU 上执行引擎对一张随机选取的 MNIST 手写数字图片PGM 格式进行分类并打印 10 个类别的概率。示例源码全部位于 sampleOnnxMNIST.cpp其类SampleOnnxMNIST只暴露两个公共方法build()构建引擎与infer()执行推理结构极其清晰非常适合作为阅读 TensorRT C API 的第一份源码。二、将 ONNX 模型转换为 TensorRT 网络2.1 初始化 ONNX ParserONNX 模型转换为 TensorRT 网络的工作由 ONNX Parser 完成。Parser 初始化时需要两个对象它将要写入的网络定义INetworkDefinition和日志对象ILogger。对应源码见 sampleOnnxMNIST.cppauto parser std::unique_ptrnvonnxparser::IParser(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));network是在build()中通过builder-createNetworkV2(...)创建的INetworkDefinitionParser 解析出的层、权重与张量关系都会写入其中sample::gLogger.getTRTLogger()返回 TensorRT 全局日志器用于报告解析过程中的错误、警告与信息。该 Parser 接口定义在 include/NvOnnxParser.h对应实现位于仓库的 parsers/onnx 目录属于 TensorRT 的开源组件之一。除源码集成外TensorRT 也支持通过nvinfer_onnxparser动态库加载 Parser。2.2 解析 ONNX 文件ONNX 模型文件通过parseFromFile传入 Parser同时传入日志级别决定哪些严重级别的日志会被上报见 sampleOnnxMNIST.cppauto parsed parser-parseFromFile(samplesCommon::locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(), static_castint(sample::gLogger.getReportableSeverity()));原 README 中给出的等价写法及失败处理逻辑为if (!parser-parseFromFile(model_file, static_castint(sample::gLogger.getReportableSeverity()))) { string msg(failed to parse onnx file); sample::gLogger-log(nvinfer1::ILogger::Severity::kERROR, msg.c_str()); exit(EXIT_FAILURE); }两个值得留意的细节模型文件定位这里没有直接使用裸路径而是经过samplesCommon::locateFile在多个候选目录中查找。该函数实现在 samples/common/common.h会依次遍历传入的目录列表若当前目录找不到还会沿父目录向上回溯最多 10 层MAX_DEPTH因此从不同工作目录运行示例都能找到mnist.onnx。找到文件后readPGMFile等数据读取函数同样依赖这套查找机制。日志级别sample::gLogger.getReportableSeverity()决定 Parser 在解析过程中输出哪些级别的日志便于定位模型算子不支持等解析问题。解析完成后mnist.onnx中的网络结构本示例包含卷积、激活、池化、矩阵乘、Scale、Shuffle 等算子详见第五节就被完整写入network之后即可进入引擎构建阶段。三、构建引擎从网络定义到可执行引擎3.1 创建 Builder、网络与配置构建引擎的第一步是创建IBuilder并传入 TensorRT 日志器用于上报构建过程中的错误、警告与信息IBuilder* builder createInferBuilder(sample::gLogger);在 sampleOnnxMNIST.cpp 的build()中接着创建网络定义与构建配置auto network std::unique_ptrnvinfer1::INetworkDefinition( builder-createNetworkV2(1U static_castuint32_t(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED))); auto config std::unique_ptrnvinfer1::IBuilderConfig(builder-createBuilderConfig());值得注意的是createNetworkV2使用了kSTRONGLY_TYPED标志——这与 README Changelog 中 2025 年 10 月的更新Migrate to strongly typed APIs相吻合。强类型网络要求所有张量显式指定数据类型如kFLOATTensorRT 不再自动推断类型这避免了隐式转换带来的歧义也是当前版本推荐的做法。3.2 可选的 Timing Cache 与 Profile Stream在正式构建前constructNetwork()中还做了两项可选配置见 sampleOnnxMNIST.cppif (mParams.timingCacheFile.size()) { timingCache samplesCommon::buildTimingCacheFromFile(sample::gLogger.getTRTLogger(), *config, mParams.timingCacheFile); } samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);Timing Cache当通过--timingCacheFile指定路径时构建前会从该文件加载已有的计时缓存ITimingCache复用之前的 kernel 选择结果从而缩短重复构建时间构建完成后若缓存文件尚不存在则会创建它。DLA 支持--useDLACore指定 DLADeep Learning Accelerator核心后会通过samplesCommon::enableDLA为支持 DLA 的层启用 DLA 执行。此外sampleOnnxMNIST.cpp 中为 Builder 设置了专用的 CUDA profile streamauto profileStream samplesCommon::makeCudaStream(); config-setProfileStream(*profileStream);3.3 序列化构建与反序列化引擎构建的核心调用是buildSerializedNetwork它直接返回序列化后的引擎二进制IHostMemory见 sampleOnnxMNIST.cppstd::unique_ptrIHostMemory plan{builder-buildSerializedNetwork(*network, *config)};这一步是 TensorRT 优化的核心Builder 会根据目标 GPU 架构执行层融合、kernel 自动调优、内存布局优化等一系列编译期工作耗时通常远大于解析阶段。构建完成后用运行时反序列化得到可直接执行的引擎对象mRuntime std::shared_ptrnvinfer1::IRuntime(createInferRuntime(sample::gLogger.getTRTLogger())); mEngine std::shared_ptrnvinfer1::ICudaEngine(mRuntime-deserializeCudaEngine(plan-data(), plan-size()));在 sampleOnnxMNIST.cpp 中build()还对网络 I/O 维度做了断言校验ASSERT(network-getNbInputs() 1); mInputDims network-getInput(0)-getDimensions(); ASSERT(mInputDims.nbDims 4); // NCHW即 1x1x28x28 ASSERT(network-getNbOutputs() 1); mOutputDims network-getOutput(0)-getDimensions(); ASSERT(mOutputDims.nbDims 2); // 10 个类别即输入为四维1x1x28x28NCHW输出为二维1x1010 个数字类别的分数。原 README 也特别提示引擎构建完成后应通过确认输出是否符合预期来验证引擎运行正确本示例的输出格式与 sampleMNIST 完全一致。四、运行推理输入预处理、执行与结果校验原 README 指出推理执行可参考 TensorRT 开发者指南中 Performing Inference In C 一节并强调对输入数据进行预处理、转换成网络接受的格式至关重要。在本示例中样本输入为 PGMportable graymap灰度图格式网络期望的输入是尺寸1x28x28、数值归一化到[0,1]区间的张量。下面结合源码拆解infer()的完整链路见 sampleOnnxMNIST.cpp。4.1 创建上下文与绑定张量地址samplesCommon::BufferManager buffers(mEngine); auto context std::unique_ptrnvinfer1::IExecutionContext(mEngine-createExecutionContext()); for (int32_t i 0, e mEngine-getNbIOTensors(); i e; i) { auto const name mEngine-getIOTensorName(i); context-setTensorAddress(name, buffers.getDeviceBuffer(name)); }BufferManager是 samples/common/buffers.h 中提供的 RAII 缓冲管理类会为每个 I/O tensor 同时分配 host 端与 device 端内存并自动维护cudaMalloc/cudaFree当前示例通过setTensorAddress显式地为每个 I/O tensor 绑定设备内存地址对应 README Changelog 中迁移到强类型 API的另一个体现替代了旧的 binding index 方式。4.2 输入预处理PGM → 浮点张量processInput()见 sampleOnnxMNIST.cpp完成输入侧的全部工作随机选数字用std::uniform_int_distributionint{0, 9}随机抽取一个 0~9 的数字mNumber读取 PGM 文件通过locateFile定位{数字}.pgm如3.pgm再用 samples/common/common.h 中的readPGMFile按二进制读取 28x28 个灰度像素PGM 文件的魔数、宽高、最大灰度值等头部信息被跳过打印 ASCII 预览用 .:-*#%[fileData[i] / 26]将灰度值映射为字符在终端以字符画形式打印输入数字归一化与反色将像素值转换为网络期望的浮点格式float* hostDataBuffer static_castfloat*(buffers.getHostBuffer(mParams.inputTensorNames[0])); for (int i 0; i inputH * inputW; i) { hostDataBuffer[i] 1.0 - float(fileData[i] / 255.0); }这里的1.0 - gray/255.0是关键细节MNIST 原图为黑底白字而模型训练时数字笔画为高值约 1.0、背景为低值约 0.0因此做了反色处理使预处理后的输入与训练分布一致。这正是 README 强调预处理必须匹配网络输入格式的具体体现。4.3 执行推理与结果回拷buffers.copyInputToDevice(); // Host → Device bool status context-executeV2(buffers.getDeviceBindings().data()); buffers.copyOutputToHost(); // Device → HostcopyInputToDevice/copyOutputToHost封装在BufferManager中samples/common/buffers.h内部按 tensor 的输入/输出模式自动选择cudaMemcpyHostToDevice或cudaMemcpyDeviceToHostcontext-executeV2在 GPU 上同步执行推理另有enqueueV3等异步接口适用于流式场景。4.4 输出校验Softmax 与置信度阈值verifyOutput()见 sampleOnnxMNIST.cpp对网络输出做后处理与校验float sum{0.0F}; for (int i 0; i outputSize; i) { output[i] exp(output[i]); // 指数化 sum output[i]; // 求和 } // ... 每个元素除以 sum 得到概率并记录最大概率及其下标 idx return idx mNumber val 0.9F;即对 10 个原始 logits 计算 Softmax得到各数字概率并打印概率条*数量表示概率大小最终要求预测类别等于真实数字mNumber且置信度val 0.9才判定推理成功。这为示例提供了一个可自动验证的通过标准也对应了运行日志末尾的 PASSED。五、本示例涉及的 TensorRT 层与算子原 README 完整列出了 sampleOnnxMNIST 网络用到的 TensorRT 层以下是逐层说明这些层来自mnist.onnx的网络结构由 ONNX Parser 转换为 TensorRT 层层类型作用在本示例中的角色Activation 层实现逐元素激活函数使用类型kRELU即 ReLU 激活Convolution 层计算 2D通道、高、宽卷积可带 bias提取手写数字的局部特征MatrixMultiply 层矩阵乘法运算实现全连接分类见下方说明Pooling 层通道内池化支持maximum、average及二者混合下采样、降低特征图分辨率Scale 层逐张量 / 逐通道 / 逐元素的仿射变换或常数幂运算归一化或缩放中间特征Shuffle 层张量的 reshape 与 transpose调整特征图维度以衔接各层其中关于全连接层需要特别说明TensorRT 的FullyConnected 层自 8.4 版本起已废弃官方建议改用MatrixMultiply层且全连接层的偏置可通过Elementwise层的SUM运算叠加——这正是本示例网络采用 MatrixMultiply 的原因也提醒读者在自建网络时优先使用非废弃 API。各层的完整接口定义可查阅 include/NvInfer.h。六、运行示例前的准备编译与数据6.1 编译示例sampleOnnxMNIST 由 CMakeLists.txt 定义构建规则它将sampleOnnxMNIST.cpp编译为可执行文件sample_onnx_mnist并链接trt_samples_commonsamples 公共库位于 samples/common与TRT_SAMPLES::tensorrt目标。编译请遵循仓库根目录 README.md 中关于构建 TensorRT 及 samples 的说明通常在完成 TensorRT 整体 CMake 配置后samples 会一并构建产物中即包含sample_onnx_mnist。6.2 准备样本数据MNIST 模型文件与测试图片需要提前就位具体步骤见 samples/README.md下载 TensorRT sample data 数据包若使用 NVIDIA NGC 容器数据通常已挂载在/usr/src/tensorrt/data下可跳过下载解压并放置到指定目录设置环境变量unzip tensorrt_sample_data_xxx.zip mkdir -p /usr/src/tensorrt/data cp -r tensorrt_sample_data_*/* /usr/src/tensorrt/data/ export TRT_DATADIR/usr/src/tensorrt/data解压后的数据目录应包含int8_api/、mnist/、resnet50/等子目录其中mnist/下就是本示例需要的mnist.onnx模型与0.pgm~9.pgm测试图片。若未指定--datadir示例默认在data/samples/mnist/与data/mnist/两个相对目录中查找见initializeSampleParamssampleOnnxMNIST.cpp。七、运行示例与完整输出解读7.1 启动命令编译并准备数据后运行./sample_onnx_mnist [-h or --help] [-d or --datadirpath to data directory] [--useDLACoreint]对应参数在源码printHelpInfo()sampleOnnxMNIST.cpp中有完整定义选项含义-h/--help显示帮助信息-d/--datadirpath指定数据目录可多次使用以添加多个目录缺省时使用data/samples/mnist/、data/mnist/--useDLACoreN为支持 DLA 的层指定 DLA 引擎取值 0 到 n-1n 为平台上 DLA 引擎数量-t/--timingCacheFilepath指定 timing cache 文件路径若文件不存在则创建这些参数在 samples/common/argsParser.h 的parseArgs中解析其中--datadir被 push 进dataDirs向量、--useDLACore解析为整数、--timingCacheFile保存为字符串最终由initializeSampleParams组装成OnnxSampleParams传给示例类。7.2 运行日志解读成功运行时的日志形如模型元数据部分来自mnist.onnx文件头由 ONNX Parser 打印 RUNNING TensorRT.sample_onnx_mnist # ./sample_onnx_mnist ---------------------------------------------------------------- Input filename: ../../../../../../data/samples/mnist/mnist.onnx ONNX IR version: 0.0.3 Opset version: 1 Producer name: CNTK Producer version: 2.4 Domain: Model version: 1 Doc string: ---------------------------------------------------------------- [I] Input: ... [I] Output: Prob 0 0.0000 Class 0: Prob 1 0.0000 Class 1: Prob 2 1.0000 Class 2: ********** Prob 3 0.0000 Class 3: ... PASSED TensorRT.sample_onnx_mnist # ./sample_onnx_mnist日志分四段模型元信息ONNX IR 版本、Opset 版本、Producer本模型由 CNTK 2.4 导出等帮助确认加载的是预期模型输入 ASCII 预览随机选中的手写数字以字符画形式打印可直观核对输入是否正确输出概率分布10 个类别的 Softmax 概率及概率条Prob 2 1.0000表示数字 2 的置信度接近 1测试结论末尾的 PASSED表示校验通过预测类别正确且置信度 0.9。若失败则会输出 FAILED并返回非零退出码。八、端到端调用链小结把整个示例串起来一次完整运行的源码级调用链为main() [sampleOnnxMNIST.cpp#L385-L419] ├─ parseArgs() [samples/common/argsParser.h#L83-L158] ├─ initializeSampleParams() [sampleOnnxMNIST.cpp#L344-L363] ├─ SampleOnnxMNIST::build() [sampleOnnxMNIST.cpp#L106-L183] │ ├─ createInferBuilder createNetworkV2(kSTRONGLY_TYPED) createBuilderConfig │ ├─ nvonnxparser::createParser() [sampleOnnxMNIST.cpp#L127-L128] │ ├─ constructNetwork() [sampleOnnxMNIST.cpp#L193-L213] │ │ ├─ parser-parseFromFile(mnist.onnx) [sampleOnnxMNIST.cpp#L197-L198] │ │ ├─ buildTimingCacheFromFile() 可选--timingCacheFile │ │ └─ enableDLA() 可选--useDLACore │ ├─ config-setProfileStream() │ ├─ builder-buildSerializedNetwork() [sampleOnnxMNIST.cpp#L150] │ └─ runtime-deserializeCudaEngine() [sampleOnnxMNIST.cpp#L168] └─ SampleOnnxMNIST::infer() [sampleOnnxMNIST.cpp#L221-L264] ├─ BufferManager(mEngine) [samples/common/buffers.h#L237] ├─ context-setTensorAddress() [sampleOnnxMNIST.cpp#L232-L236] ├─ processInput() [sampleOnnxMNIST.cpp#L269-L299] │ ├─ locateFile readPGMFile [samples/common/common.h#L526-L593] │ └─ 归一化/反色 → host buffer ├─ buffers.copyInputToDevice() → executeV2() → copyOutputToHost() └─ verifyOutput()Softmax 置信度校验 [sampleOnnxMNIST.cpp#L306-L339]九、延伸学习与注意事项延伸学习仓库内资源本示例使用的 Parser 头文件为 include/NvOnnxParser.h解析器开源实现位于 parsers/onnx若想进一步学习 ONNX 导入可阅读 samples/python/introductory_parser_samples/README.mdPython 版 Parser 入门以及 engine_refit_onnx_bidaf基于 Parser 构建引擎后的 refit动态形状的 ONNX 示例见 sampleDynamicReshapeI/O 格式指定见 sampleIOFormats命名维度见 sampleNamedDimensions定时缓存timing cache的深入用法可参考 sampleEditableTimingCache命令行工具 trtexec 无需写代码即可完成 ONNX → 引擎 → 推理的同一流程适合快速验证模型。注意事项预处理必须与训练一致PGM 像素需反色并归一化到[0,1]1.0 - gray/255.0否则分类精度会显著下降强类型网络本示例使用kSTRONGLY_TYPED创建网络所有张量必须显式指定数据类型数据目录务必确保mnist.onnx与0.pgm~9.pgm位于--datadir指定目录或默认的data/samples/mnist/、data/mnist/下环境依赖示例需要 NVIDIA GPU、CUDA 环境以及已构建的 TensorRT 库与 samples 公共库已知问题原 README 声明该示例暂无已知问题Known issues: none。sampleOnnxMNIST 虽然短小却完整覆盖了ONNX 解析 → 引擎构建 → 数据预处理 → 推理执行 → 结果校验这条 TensorRT 使用主链路是理解 TensorRT 工作流的最佳起点。建议读者在跑通示例后对照 sampleOnnxMNIST.cpp 逐行阅读build()与infer()即可快速建立对 TensorRT C API 的整体认知。【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考