PVN3D TensorRT Engine 转换与测试记录
1. 目的本文单独记录当前pvn3d-dev容器内从 ONNX 到 TensorRT engine 的转换流程、测试方法以及这一阶段需要注意的约束。这篇文档只关注两件事如何把当前已经导出的 ONNX 子图转换成 TensorRT engine如何验证 engine 是否构建成功以及这条链路当前有哪些边界不在本文展开的内容TensorRT 安装ONNX 导出问题本身PointNet2 TensorRT plugin相关文档pvn3d_onnx_export_and_test.mdpvn3d_tensorrt_env_setup_guide.mdtensorrt_install_in_pvn3d_dev.md2. 当前容器环境本文基于当前实测环境编写容器pvn3d-dev仓库路径/workspace/workflow/self/PVN3DConda 环境pvn3dPython3.8.20PyTorch1.10.0cu113ONNX1.14.1ONNX Runtime1.16.3TensorRT Python8.6.1trtexec/opt/tensorrt/TensorRT-8.6.1.6/bin/trtexectrtexec --help可执行版本标识为TensorRT v8601当前环境已经满足 engine 构建前提。3. 当前 TensorRT 转换目标当前不是把整个 PVN3D 一次性转换成单个 TensorRT engine。当前转换目标只有两个rgb_backbone.onnx - rgb_backbone.enginefusion_head.onnx - fusion_head.enginePointNet2 仍然保留原生 CUDA 路径。所以当前实际部署结构仍然是rgb_backbone.enginePointNet2 Native CUDAfusion_head.engine这条路线延续了 ONNX 拆分阶段的同一原则先把标准算子图转换成 TensorRT把 PointNet2 自定义 CUDA 算子留在现有路径。4. 当前已有 ONNX 产物当前仓库里已经存在两组 ONNXdeploy/models/onnx/*.onnxdeploy/models/onnx_ape/*.onnx对 LINEMODape的测试应优先使用deploy/models/onnx_ape/rgb_backbone.onnxdeploy/models/onnx_ape/fusion_head.onnx原因onnx_ape是用ape_pvn3d_best.pth.tar导出的num_classes2与 LINEMODape的单类配置一致不要拿YCB 22 类的 ONNX 去构建 LINEMODape的 engine。这里也补一条边界说明当前 LINEMOD 单物体流程里num_classes通常固定为2所以如果后续从ape换到其他 LINEMOD 物体通常不需要修改num_classes例如apenum_classes2cannum_classes2drillernum_classes2但即便如此下面这些仍然不能混用checkpointONNX 目录engine 目录测试脚本里的--cls原因不是类别数不同而是虽然都属于“单前景类”这个前景类对应的具体物体不同权重不同mesh 不同keypoints 不同center 不同pose 求解目标也不同所以对 LINEMOD 更准确的理解是num_classes2往往不变但ape、can、driller仍应各自导出各自的 ONNX 和 engine5. 进入环境dockerexec-itpvn3d-devbashsource/opt/conda/etc/profile.d/conda.sh conda activate pvn3dcd/workspace/workflow/self/PVN3D建议先确认whichtrtexec python -PY import tensorrt as trt print(TensorRT:, trt.__version__) PY6. 当前 engine 构建脚本当前仓库使用build_trt_engine.py这个脚本本质上是对trtexec的封装。它当前负责检查trtexec是否存在在需要时按目标子图自动拼接输入 shape生成并执行trtexec命令输出.engine当前支持目标rgb_backbonefusion_head7. 构建前的关键约束7.0 先区分静态 ONNX 和动态 ONNX当前deploy/models/onnx_ape/*.onnx已经确认是静态输入模型rgb_backbone.onnxrgb: [1, 3, 480, 624]fusion_head.onnxout_rgb: [1, 128, 480, 624]choose: [1, 1, 4096]pcld_emb: [1, 128, 4096]这意味着当前构建 engine 时不应该再给trtexec传--minShapes--optShapes--maxShapes如果对静态 ONNX 强行传这些参数TensorRT 会直接报错。7.1 分辨率必须与导出 ONNX 一致当前已经验证通过的 ONNX 导出尺寸是height480width624所以构建 TensorRT engine 时也必须使用同样的输入尺寸。不要再写成--height480--width640当前推荐值是--height480--width6247.2 点数必须与导出 ONNX 一致当前导出使用的是--num-points4096所以fusion_head.engine构建时也必须保持--num-points40967.3 ONNX 和 engine 的类别语义必须一致如果构建 LINEMODape的 engine优先使用deploy/models/onnx_ape不要把deploy/models/onnx和deploy/models/onnx_ape混着用。8. 推荐构建命令8.1 先做 dry-run先只打印命令不实际构建python deploy/scripts/build_trt_engine.py\--onnx-dir deploy/models/onnx_ape\--engine-dir deploy/models/engine_ape\--targetall\--fp16\--height480\--width624\--num-points4096\--dry-run这样做的目的先确认onnx-dir是否正确先确认静态 ONNX 没有被错误地附加动态 profile shape先确认脚本实际拼出来的trtexec命令是否符合预期8.2 实际构建确认 dry-run 没问题后再执行python deploy/scripts/build_trt_engine.py\--onnx-dir deploy/models/onnx_ape\--engine-dir deploy/models/engine_ape\--targetall\--fp16\--height480\--width624\--num-points40968.3 只构建rgb_backbonepython deploy/scripts/build_trt_engine.py\--onnx-dir deploy/models/onnx_ape\--engine-dir deploy/models/engine_ape\--targetrgb_backbone\--fp16\--height480\--width624\--num-points40968.4 只构建fusion_headpython deploy/scripts/build_trt_engine.py\--onnx-dir deploy/models/onnx_ape\--engine-dir deploy/models/engine_ape\--targetfusion_head\--fp16\--height480\--width624\--num-points40969. 当前脚本对 shape 的处理逻辑9.1 静态 ONNX如果 ONNX 输入维度都是固定值脚本现在会不传--minShapes不传--optShapes不传--maxShapes这是当前onnx_ape的实际情况。例如rgb_backbone的 dry-run 结果现在是trtexec--onnxdeploy/models/onnx_ape/rgb_backbone.onnx\--saveEnginedeploy/models/engine_ape/rgb_backbone.engine\--memPoolSizeworkspace:4096\--skipInference\--fp169.2 动态 ONNX只有当 ONNX 本身包含动态维度时脚本才会补--minShapes--optShapes--maxShapes当前这套onnx_ape不属于这种情况。9.3 如果模型是动态 ONNX目标 shape 约定如下脚本会传给trtexecrgb:1x3x480x624对应input:rgbshape:[1, 3, 480, 624]9.4fusion_head脚本会传给trtexecout_rgb:1x128x480x624 choose:1x1x4096 pcld_emb:1x128x4096也就是out_rgb:[1, 128, 480, 624]choose:[1, 1, 4096]pcld_emb:[1, 128, 4096]这与当前导出脚本的输入约定保持一致。10. 构建后检查构建成功后先检查输出目录finddeploy/models/engine_ape-maxdepth1-typef|sort预期文件deploy/models/engine_ape/rgb_backbone.enginedeploy/models/engine_ape/fusion_head.engine11. engine 基础测试11.1 用trtexec加载 engine先单独测rgb_backbone.enginetrtexec\--loadEnginedeploy/models/engine_ape/rgb_backbone.engine\--skipInference然后测fusion_head.enginetrtexec\--loadEnginedeploy/models/engine_ape/fusion_head.engine\--skipInference这里的目的不是测性能而是确认engine 能被 TensorRT 正常反序列化没有明显的版本或 shape 错误11.2 看帮助确认支持项当前容器里的trtexec --help已经验证通过说明 CLI 可用。如果后续要补性能测试可以继续使用--dumpProfile--dumpLayerInfo--shapes...12. 与混合推理测试的关系这篇文档只负责“ONNX - engine”的构建与最小验证。如果你要验证完整推理链应结合run_linemod_hybrid_ort.pyrun_linemod_hybrid_trt.py当前更合理的推进顺序是ONNX 导出成功ONNX Runtime 能加载 ONNXtrtexec能成功构建 enginetrtexec --loadEngine能成功加载 engine再把 ONNX Runtime 版混合推理逐步替换成 TensorRT engine 版混合推理13. TensorRT Python 混合测试13.1 当前是否可以写 TensorRT Python 测试脚本可以。当前容器里已经确认tensorrt8.6.1torch1.10.0cu113没有pycuda没有cupy因此当前 Python 测试脚本采用的是TensorRT Python API 负责反序列化 engine 和执行torch.cuda负责分配 GPU tensor通过tensor.data_ptr()把地址传给 TensorRT这条方式不依赖pycuda。13.2 脚本位置脚本放在run_linemod_hybrid_trt.py脚本做的事情读取一帧 LINEMOD 样本把 RGB 从640裁成624把点数从12288重采样到4096用 TensorRT engine 跑rgb_backbone用原生 PointNet2 CUDA 跑点云特征用 TensorRT engine 跑fusion_head可选对比整模型 PyTorch 输出调用 LINEMOD pose solver输出 pose、ADD、ADD-S13.3 运行前提这个脚本依赖两个 engine 都已经存在deploy/models/engine_ape/rgb_backbone.enginedeploy/models/engine_ape/fusion_head.engine少一个都不成立。13.4 运行命令python onnx_test/run_linemod_hybrid_trt.py\--clsape\--checkpointweights/ape_pvn3d_best.pth.tar\--rgb-engine deploy/models/engine_ape/rgb_backbone.engine\--fusion-engine deploy/models/engine_ape/fusion_head.engine\--sample-index0\--num-points4096\--height480\--width624\--crop-left8\--outputonnx_test/outputs/linemod_ape_hybrid_trt_summary.json这里可以继续使用相对路径参数不需要手工改成绝对路径。当前脚本已经会在启动时自动把下面这些参数解析成仓库根目录下的绝对路径--checkpoint--rgb-engine--fusion-engine--output这样做的原因是脚本内部为了兼容 LINEMOD 的 legacy 工具代码会切工作目录到pvn3d/如果不先把这些路径转成绝对路径weights/...、deploy/...这类相对路径会被带偏所以当前推荐用法仍然是文档里的相对路径写法。13.5 输出内容脚本会输出并保存一份 JSON总结以下信息使用的 checkpoint使用的 engine 文件输入 shape点云 shape预测 poseADDADD-S如果未加--skip-torch-compare还会给出混合推理和整模型 PyTorch 结果的误差统计14. 当前已确认的问题14.1build_trt_engine.py的默认宽度还是640当前脚本里的默认参数仍是parser.add_argument(--width,typeint,default640,help图像宽。)这和当前已经验证通过的 ONNX 尺寸480x624不一致。这不表示脚本不能用但表示运行时必须显式传--width 624不能依赖默认值14.2 静态 ONNX 不能传显式 shape profile当前已经实际遇到过下面的失败Static model does not take explicit shapes since the shape of inference tensors will be determined by the model itself触发原因rgb_backbone.onnx是静态模型但构建时仍传了--minShapes/--optShapes/--maxShapes这不是 ONNX 损坏也不是 TensorRT 安装有问题而是trtexec参数和模型类型不匹配。解决已经更新 build_trt_engine.py现在脚本会自动检查 ONNX 是否是静态输入对静态 ONNX不再传动态 profile shape14.3trtexec --version不会只打印版本在当前 TensorRT 8.6.1 环境里trtexec--version会先打印版本标识TensorRT v8601但后面仍可能继续进入参数解析并报缺少模型。这不表示 TensorRT 有问题只是这个 CLI 的行为方式如此。14.4 TRT Python 测试脚本的路径也会受chdir影响当前run_linemod_hybrid_trt.py内部会切到pvn3d/工作目录以兼容 LINEMOD 旧工具代码里的相对路径假设。如果脚本不先处理外部传入路径就会出现这类问题weights/ape_pvn3d_best.pth.tar被错误解析deploy/models/engine_ape/...被错误解析当前脚本已经修复先把外部传入路径转成仓库根目录下的绝对路径再切工作目录因此现在文档里的相对路径命令是可直接使用的。15. 当前处理建议15.1 构建 engine 时显式传入480x624不要依赖脚本默认值。15.2 对 LINEMODape使用onnx_ape不要混用 YCB 的 ONNX。15.3 先 dry-run 再实际构建这样最容易发现 shape 和路径错误。15.4 对静态 ONNX不要手工再补动态 profile shape如果脚本已经识别为静态 ONNX就不要手工往命令里加--minShapes--optShapes--maxShapes15.5 先做最小 engine 验证再做整链替换当前阶段优先确认ONNX 正常engine 正常TensorRT 反序列化正常然后再继续做混合推理替换。15.6 有两个 engine 以后再跑 TRT Python 混合测试当前 TRT Python 测试脚本依赖rgb_backbone.enginefusion_head.engine少一个都不成立。15.7 直接按文档命令运行不需要手工改绝对路径当前文档里的 TRT Python 测试命令已经与脚本实现保持一致。继续使用weights/...deploy/models/...onnx_test/outputs/...这种相对路径写法即可。16. 当前建议的测试顺序建议按下面顺序执行确认deploy/models/onnx_ape/*.onnx存在python deploy/scripts/build_trt_engine.py ... --dry-runpython deploy/scripts/build_trt_engine.py ...trtexec --loadEngine... --skipInferencepython onnx_test/run_linemod_hybrid_trt.py ...再考虑把 ORT 混合推理替换成 TRT 混合推理