NVIDIA TensorRT 中带自定义层 ONNX 模型推理实战PackNet 自监督单目深度估计示例精讲【免费下载链接】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本文基于 NVIDIA TensorRT 开源仓库中的 samplePackNet 示例samples/python/onnx_packnet系统讲解如何将带有自定义层的 PyTorch 网络转换为 ONNX再通过 ONNX Parser 与 ONNX-GraphSurgeon 图变换接入 TensorRT 推理。读完本文你将掌握「PyTorch → ONNX → 图结构精简 → 自定义插件替换 → trtexec 推理」的完整实战链路并能独立迁移到其他包含 GroupNorm、Upsample、Pad 等难解析算子的模型中。一、示例概述PackNet 与 samplePackNetPackNet 是一个用于自动驾驶场景的自监督单目深度估计网络self-supervised monocular depth estimation network其核心思路是通过 3D 打包/解包packing/unpacking操作在多个尺度上保留空间信息实现高精度深度预测。该网络因大量使用 Group Normalization、Upsample、Pad 等算子恰好是检验 TensorRT「自定义层plugin 图优化」能力的典型模型。本示例 samplePackNet 的完整技术路径是将 PyTorch 训练好的 PackNet01 网络导出为 ONNX 图使用ONNX-GraphSurgeonONNX-GS对 ONNX 图做结构化精简折叠 Pad 子图、折叠 Upsample 子图、把 Group Normalization 子图替换为单个插件节点由 TensorRT 内置的ONNX Parser解析处理后的 ONNX 图通过REGISTER_TENSORRT_PLUGINAPI 让自定义层GroupNormalizationPlugin在 TensorRT 中自动注册无需手动加载插件库使用trtexec工具对随机输入图像执行端到端推理。整个示例的核心代码仅三个文件convert_to_onnx.py导出与主流程、post_processing.pyONNX-GS 图变换、requirements.txt依赖清单外加一个用于预下载 packnet 源码包的 download.yml。二、工作原理ONNX 解析与自定义层自动注册原文档明确描述了该示例演示的两个关键技术点这也是把任意 PyTorch 模型接入 TensorRT 的通用方法论1. 自定义层plugin的自动注册ONNX 图里的 Group Normalization 在 PyTorch 导出时会展开成一长串子图InstanceNormalization Reshape Scale/Bias 等TensorRT 解析器无法高效识别。示例的做法是先用 ONNX-GS 把子图折叠成单个GroupNormalizationPlugin节点再由插件注册机制接管。注册由REGISTER_TENSORRT_PLUGIN宏完成——在 groupNormalizationPlugin.cpp 中可见REGISTER_TENSORRT_PLUGIN(GroupNormalizationPluginCreator);该宏将GroupNormalizationPluginCreator实例自动加入 TensorRT 插件注册表因此在用trtexec --onnxmodel.onnx或 Python API 构建引擎时解析器遇到名为GroupNormalizationPlugin的节点会自动按插件名、版本号1与命名空间查找并实例化无需额外dlopen或手动initLibNvInferPlugins。2. ONNX-GraphSurgeon 图变换ONNX-GS 负责在解析前对 ONNX 图做「减肥」删除推理阶段无用的中间节点、将冗余子图折叠为单节点。具体到 PackNet变换目标是 Group Normalization、Upsample 和 Pad 三类子图细节见下文第四节。三、环境准备与依赖安装按原文档运行示例前需升级 pip 并安装依赖pip3 install --upgrade pip pip3 install -r requirements.txtrequirements.txt 的实际内容如下onnx1.18.0 --extra-index-url https://pypi.ngc.nvidia.com onnx-graphsurgeon0.3.20 torch torchvision pyyaml6.0.3 requests2.32.4 tqdm4.66.4 numpy1.26.4几点补充说明onnx-graphsurgeonONNX-GS是执行图变换的关键库示例要求不低于 0.3.20它与 TensorRT 一同发布也可通过 NGC 的 PyPI 源--extra-index-url指定的https://pypi.ngc.nvidia.com安装onnx1.18.0与示例支持的 ONNX opset 11 匹配见下文参数说明需要同时具备一个可用的 TensorRT Python 包提供trtexec与 ONNX Parser原文档特别提醒在 PowerPC 系统上需要改用 IBM PowerAI 渠道手动安装 PyTorch不能直接走 pip。四、准备 PackNet 模型仓库PackNet 本体来自 TRI-ML 的 packnet-sfm 开源实现。原文档给出的准备步骤为git clone https://github.com/TRI-ML/packnet-sfm.git packnet-sfm pushd packnet-sfm git checkout tags/v0.1.2 popd export PYTHONPATH$PWD/packnet-sfm # Windows 下为: set PYTHONPATH%cd%\packnet-sfm必须检出v0.1.2标签因为示例代码针对该版本的网络结构与算子导出行为编写后续图变换中多处torch.__version__分支即是为兼容不同导出格式将仓库根目录加入PYTHONPATH使 convert_to_onnx.py 中的from packnet_sfm.networks.depth.PackNet01 import PackNet01可以正常导入仓库同时提供了 download.yml其中记录了 packnet-sfm v0.1.2 源码包的下载地址与 MD5 校验和7a73db591d3955ccf407910cd928d9c0可供离线环境核对下载产物完整性。五、PyTorch → ONNX 转换与图后处理5.1 转换命令与命令行参数原文档给出的核心命令是python3 convert_to_onnx.py --output model.onnx从 convert_to_onnx.py 的argparse定义可看到完整参数表参数简写类型默认值说明--output-ostrmodel.onnx生成的 ONNX 模型保存路径--opset-opint11导出使用的 ONNX opset 版本--verbose-vflag关闭开启torch.onnx.export的详细日志运行convert_to_onnx.py -h或--help可随时查看上述完整说明。示例默认使用 opset 11这决定了后处理脚本中 Upsample/Resize、Pad 子图的折叠分支选择详见第六节。5.2 主流程解析convert_to_onnx.py 的主流程分两步main()中依次调用第一步build_packnet构建并导出网络input_pyt torch.randn((1, 3, 192, 640), requires_gradFalse) model_pyt PackNet01(version1A) torch.onnx.export(model_pyt, input_pyt, model_file, verboseargs.verbose, opset_versionargs.opset)输入为(1, 3, 192, 640)的随机张量单张 192×640 的三通道图像即 PackNet 的标准输入分辨率网络选用PackNet01(version1A)变体固定输入尺寸导出得到的 ONNX 图在 TensorRT 中按静态 shape 构建引擎。第二步post_process_packnetONNX-GS 图后处理graph gs.import_onnx(onnx.load(model_file)) if opset 11: graph process_pad_nodes(graph) if torch.__version__ 1.5.0: graph process_upsample_nodes(graph, opset) graph process_groupnorm_nodes(graph) graph.cleanup().toposort() onnx.save_model(gs.export_onnx(graph), model_file)流程要点用gs.import_onnx载入图依次执行 Pad 折叠、Upsample 折叠仅针对旧版 torch 导出的子图结构、GroupNorm 折叠graph.cleanup().toposort()删除不再被引用的孤立节点并对图做拓扑排序保证输出图合法、精简最后用gs.export_onnxonnx.save_model写回model.onnx。六、三种 ONNX 图变换的源码级精讲post_processing.py 是本示例最核心的部分它针对 PyTorch 导出 ONNX 时的三个赘余子图分别做折叠。以下逐一拆解。6.1 Pad 子图折叠process_pad_nodesPyTorchopset 11 下导出的 Pad 层并不是单个 Pad 节点而是一个由Constant → Shape → Mul/Gather → Sub → ConstantOfShape → Concat → Reshape → Slice → Transpose → Reshape组成的庞杂子图最终才进入真正的Pad节点源文件顶部注释给出了完整拓扑。这样的子图在推理阶段几乎全部是常量折叠即可消除的工作。process_pad_nodes找到所有Pad节点调用fold_pad_inputs把垫多少像素的信息直接从子图中提取出来改写成[0]*8的 4D 双面 padding 常量onnx_pad_values[j]/onnx_pad_values[j4]分别对应每个维度的前后 padding并用gs.Constant直接替换 Pad 的第二个输入从而把整棵子图压缩为Input → Pad → Conv脚本中还针对不同 torch 版本做了取值路径分支1.5.0、2.0.0、2.0.0保证在不同 PyTorch 导出格式下都能正确取到 padding 值。代码注释明确假设输入为 4D 张量迁移到其他网络时若张量维数不同需要相应调整。6.2 Upsample 子图折叠process_upsample_nodes旧版 PyTorch 导出的上采样层同样是一棵复杂子图Mul → Cast → Floor → Unsqueeze → Concat → Cast → Div → Concat → Upsample。fold_upsample_inputs的目标是把它折叠为Conv scale_factor | / Upsample | ReLU实现上分两种情况opset 9从上采样输入子图的Mul中取出 scale factor构造[1.0, 1.0, scale, scale]的 scales 常量张量替换 Upsample 的最后一个输入opset 11本示例默认此时 PyTorch 导出的算子名是Resize。脚本先把Resize节点改名回Upsample再将其 4 个输入收缩为 2 个数据张量 常量 scales 张量本质上是把Resize 的 sizes 输入这种解析器不友好的形式改写为 TensorRT 解析器能直接消费的Upsample scales形式。注意原脚本中process_upsample_nodes仅在torch.__version__ 1.5.0时调用——这是历史兼容分支新版本 PyTorch 导出的上采样结构不同无需此折叠该判断逻辑在 convert_to_onnx.py。6.3 GroupNorm 子图折叠process_groupnorm_nodes与插件节点注入PyTorch 会把 Group Normalization 导出为Conv → Reshape → InstanceNormalization → Reshape → Mul(scale) → Add(bias) → ReLUprocess_groupnorm_nodes遍历所有InstanceNormalization节点通过convert_to_groupnorm将其整体替换为单个GroupNormalizationPlugin节点得到Conv → GroupNorm → ReLU替换的关键步骤对应 post_processing.py收集插件属性retrieve_attrsnum_groups取 Reshape 形状常量第二个元素即分组数eps取 InstanceNormalization 节点的epsilon属性plugin_version 1与plugin_namespace 这两项是解析器查找插件时的默认值脚本特意显式写出以便说明——即使省略解析器也会按默认值匹配创建插件节点gs.Node(opGroupNormalizationPlugin, attrsattrs)注意节点名必须与插件getPluginType()返回的GroupNormalizationPlugin完全一致见 groupNormalizationPlugin.cpp重接图拓扑把 Conv 的输出张量改为指向插件节点把 Add 的输出张量改为由插件节点输出实现从中间插入补充 scale / bias 输入把原Mulscale和Addbias的常量输入追加为插件的第 2、3 个输入——这与插件的三输入接口input scale bias严格对应。七、GroupNormalizationPlugin 插件实现剖析插件本体位于 plugin/groupNormalizationPlugin这是仓库plugin/目录下众多官方插件之一也是本示例能在 TensorRT 中正确运行 Custom Layer 的底层支撑。7.1 插件类结构groupNormalizationPlugin.h 定义GroupNormalizationPlugin它继承nvinfer1::IPluginV2DynamicExt属于动态 shape 扩展插件构造函数接收epsilon与nbGroups两个参数groupNormalizationPlugin.cpp 中还会用PLUGIN_VALIDATE校验mEpsilon 0与mNbGroups 0getOutputDimensions返回与输入相同的维度输入、输出形状一致符合归一化算子语义nbInputs 3校验三个输入内部借助 cuDNN 句柄执行归一化计算attachToContext中创建 cuDNN wrapper 与张量描述符并使用共享的CudaBindfloat缓存 scale/bias 缓冲serialize/deserialize_value按epsilon 在前、nbGroups 在后的顺序序列化参数保证引擎文件可跨进程加载配套的GroupNormalizationPluginCreator实现插件工厂接口createPlugin/deserializePlugin/getFieldNames并通过第 41 行的REGISTER_TENSORRT_PLUGIN(GroupNormalizationPluginCreator);完成自动注册。7.2 插件注册信息与解析约束仓库为插件提供了机器可读的注册描述文件 GroupNormalizationPlugin_PluginConfig.yaml其中关键约束与本示例的图变换完全咬合配置项值含义nameGroupNormalizationPlugin插件名与 ONNX 节点名一致interfaceIPluginV2DynamicExt插件接口类型versions.1.inputsinput, scale, bias3 个输入顺序与 post_processing.py 追加顺序一致input_dimsinput: 4 维scale/bias: 1 维输入为 4D 张量scale/bias 为 1D 常量input_dim_constraintsinput_1 MULTIPLE_OF num_groups_0、scale_0 input_1、bias_0 scale_0通道数须为分组数整数倍且 scale/bias 长度等于通道数attributeseps(float32)、num_groups(int32)属性名与类型与retrieve_attrs产出一致supported_input_typesfloat32 组合该配置下支持 float32 输入此外该 YAML 还给出了abs_tol/rel_tol1e-2 的数值容差以及eps/num_groups的参考测试配置如eps0.0001, num_groups1等这些信息一方面约束了解析与量化行为另一方面也是插件正确性测试golden reference的依据。八、使用 trtexec 执行 TensorRT 推理ONNX 图生成后即可用 TensorRT 自带的trtexec工具位于 TensorRT 安装包的bin目录对随机输入图像执行推理trtexec --onnxmodel.onnxtrtexec会自动完成以下工作加载model.onnx→ ONNX Parser 解析遇到GroupNormalizationPlugin节点时通过注册表自动找到插件并创建→ 构建 engine → 用随机输入数据执行推理并打印性能统计耗时、吞吐等。常用扩展用法trtexec本身支持丰富的命令行选项可用trtexec --help查看完整列表# 指定 FP16 精度与输入形状 trtexec --onnxmodel.onnx --fp16 --shapesinput:1x3x192x640 # 将构建好的 engine 保存为文件便于重复加载 trtexec --onnxmodel.onnx --saveEnginemodel.engine # 只构建不推理用于快速验证解析与插件注册是否成功 trtexec --onnxmodel.onnx --buildOnly需要说明的适用前提trtexec与 ONNX Parser 属于 TensorRT 发行包组件本仓库为 TensorRT 开源部分包含 ONNX Parser 相关头文件与文档见 NvOnnxParser.h使用--fp16等精度选项时需硬件支持Ampere 及以上架构对 FP16 有良好支持并且本示例插件配置中supported_input_types明确的是 float32 组合混合精度下请以实际构建结果为准。九、变更记录与已知问题原文档附带的 Changelog 记录了本示例的演进历史2025 年 8 月移除对 Python 3.10 的支持2023 年 8 月ONNX 版本支持更新到 1.14.0并移除对 Python 3.8 的支持2021 年 8 月适配最新 torch 版本2020 年 6 月示例首次发布。当前仓库中 requirements.txt 已将onnx固定为 1.18.0、numpy固定为 1.26.4与上述 Python 版本要求保持一致。Known issues原文档声明本示例当前无已知问题。十、延伸学习与可复用资产若想继续深入推荐结合以下仓库内资源示例代码convert_to_onnx.py 与 post_processing.py 是整个流程的完整实现可直接作为PyTorch 模型带自定义算子接入 TensorRT的模板插件实现plugin/groupNormalizationPlugin/groupNormalizationPlugin.cpp 与 groupNormalizationPlugin.h以及注册约束文件 GroupNormalizationPlugin_PluginConfig.yaml展示了自定义层从图变换命名到注册表自动发现的完整约定ONNX 解析器接口NvOnnxParser.h 定义了createParser等解析入口可进一步了解 ONNX Parser 如何消费插件节点trtexec 源码samples/trtexec/trtexec.cpp 提供了--onnx等命令行选项的完整实现细节PackNet 算法背景可查阅 TRI-ML 发布的论文《3D Packing for Self-Supervised Monocular Depth Estimation》与 packnet-sfm 开源实现v0.1.2标签下载源与校验和见 download.yml以理解网络结构为何需要如此多的 GroupNorm / Upsample / Pad 算子TensorRT 官方文档ONNX Parser 的 Python API 说明、TensorRT 开发者指南中使用 Python 导入模型章节以及 TensorRT 示例支持指南均对理解本示例有直接帮助。总结samplePackNet 用最小的代码量完整演示了 TensorRT 处理带自定义层 ONNX 模型的标准姿势先用 ONNX-GraphSurgeon 把 PyTorch 导出的冗余子图Pad / Upsample / GroupNorm折叠成解析器可消费的紧凑结构再把自定义算子替换为通过REGISTER_TENSORRT_PLUGIN自动注册的插件节点最后由 ONNX Parser 与 trtexec 无缝完成构建与推理。这套「导出 → 图变换 → 插件注册 → 解析推理」的流程具有极强的可迁移性是你在生产环境中接入任何含自定义算子的 PyTorch 模型时都可以直接复用的工程范式。【免费下载链接】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），仅供参考
