1. 文档目的本文专门说明当前仓库里deploy/model_wrappers.py、onnx_test/run_linemod_hybrid_ort.py、onnx_test/run_linemod_hybrid_trt.py三者之间是怎么配合工作的。本文讨论的是部署测试链不讨论训练。当前仓库的真实测试结构不是整个 PVN3D 一次性转成一个 ONNX整个 PVN3D 一次性转成一个 TensorRT engine而是用model_wrappers.py把完整 PVN3D 拆成三个阶段把能稳定导出的两段导出为ONNX再把 ONNX 转成TensorRT enginePointNet2 保留原生PyTorch CUDA Extension用 ORT 或 TRT 脚本把整条链重新接起来做 LINEMOD 测试这也是当前目录中以下三份代码同时存在的原因model_wrappers.pyrun_linemod_hybrid_ort.pyrun_linemod_hybrid_trt.py2. 当前测试链的核心结论当前可行的测试主线是rgb_backbone走ONNX Runtime或TensorRTPointNet2继续走原生PyTorch CUDA Extensionfusion_head走ONNX Runtime或TensorRTpose solver继续走仓库原有的 LINEMOD 求解逻辑因此当前存在两条混合测试链ORT 版rgb_backbone.onnx - PointNet2 Native CUDA - fusion_head.onnxTRT 版rgb_backbone.engine - PointNet2 Native CUDA - fusion_head.engine这两条链验证的是同一套部署拆分边界只是中间执行后端不同。3. 为什么必须拆成三段3.1 根因在 PointNet2PVN3D 里最难直接导出的部分不是 CNN也不是最终预测头而是 PointNet2。PointNet2 当前依赖本地编译的 CUDA 扩展与自定义算子这类算子不属于标准 ONNX 算子集合也不是 TensorRT 当前可以直接解析的通用层。所以当前阶段不能指望“整网直接导 ONNX”“整网直接转 TRT”3.2 当前拆分边界当前拆分边界固定为三段rgb - out_rgb, rgb_segcld_rgb_nrm - pcld_embout_rgb, choose, pcld_emb - pred_kp_of, pred_rgbd_seg, pred_ctr_of其中第 1 段可以导成rgb_backbone.onnx/rgb_backbone.engine第 2 段只能继续保留原生 CUDA第 3 段可以导成fusion_head.onnx/fusion_head.engine这不是为了代码整洁而是为了先把能稳定落地的部分落地。4. 产物之间的关系可以把当前产物链看成checkpointmodel_wrappers.pyONNXTensorRT engineORT/TRT 混合测试对应关系如下weights/*.pth.tar由 model_wrappers.py 加载wrapper 拆分后导出deploy/models/onnx_ape/rgb_backbone.onnxdeploy/models/onnx_ape/fusion_head.onnxONNX 再构建成deploy/models/engine_ape/rgb_backbone.enginedeploy/models/engine_ape/fusion_head.engine最终由测试脚本把这些产物与 PointNet2 原生 CUDA 拼回完整推理链5.model_wrappers.py代码解析文件model_wrappers.py这个文件不负责导出也不负责最终测试。它负责定义部署边界。5.1load_checkpoint职责归一化 checkpoint 路径加载 checkpoint提取state_dict去掉可能存在的module.前缀这里还有一个容器里已经遇到过的兼容性点ycb_pvn3d_best.pth.tar可以直接用torch.loadape_pvn3d_best.pth.tar可能是pickle格式所以当前实现是先尝试torch.load如果遇到Invalid magic number自动回退到pickle.load这一步解决的是“同样叫.pth.tar但底层保存格式不完全一致”的问题。5.2build_full_model职责按num_classes、num_points构建完整 PVN3D挂载 checkpoint移到目标设备切到eval()后续所有 wrapper 都建立在这一步生成的完整模型之上。5.3RGBBackboneWrapper输入rgb输出out_rgbrgb_seg对应代码逻辑很直接out_rgb,rgb_segself.cnn(rgb)returnout_rgb,rgb_seg这一段的特点是算子标准、结构清晰所以适合先导出。5.4PointNet2NativeWrapper输入pointcloud输出pcld_emb代码本身也很直接returnself.pointnet2(pointcloud)但它背后的实现依赖 CUDA 扩展这就是为什么它不能直接并入当前 ONNX / TRT 主线。5.5FusionHeadWrapper输入out_rgbchoosepcld_emb输出pred_kp_ofpred_rgbd_segpred_ctr_of这一段做的事情有三步把out_rgb拉平并按照choose抽取像素特征和pcld_emb做rgbd_feat融合通过三个 head 输出关键点偏移、分割和中心点偏移这里最关键的输入不是rgb而是已经对齐好的out_rgbchoosepcld_emb这也是 ORT/TRT 测试脚本都必须显式处理choose的原因。5.6build_split_wrappers这是整个部署测试链的工厂函数。它一次性返回full_modelrgb_backbonepointnet2_nativefusion_headgather_rgb_feature当前 ORT 和 TRT 两个测试脚本都依赖这个统一约定。6. 当前tar权重的两个作用当前混合测试链里weights/*.pth.tar不是只起一个作用而是同时承担两件事。6.1 给PointNet2 Native CUDA提供参数当前 PointNet2 还没有转成 ONNX也没有转成 TensorRT。所以无论是 ORT 测试脚本还是 TRT 测试脚本都必须先从tar权重构建完整 PVN3D再从里面取出pointnet2_native这部分权重的作用是让原生PyTorch CUDA Extension路径能够正常前向把cld_rgb_nrm转成pcld_emb作为 ONNX / TRT 两段子图之间的桥接段如果没有这部分权重当前混合链就缺中间这一段ORT 和 TRT 都接不起来。6.2 给full_model提供基线用于精度对比当前两个测试脚本里还会用同一个tar权重再构建full_model这个full_model不是当前部署链实际执行的一部分而是原始完整 PyTorch 模型基线。它的作用是把原始完整模型输出和以下结果做对比ORT 混合推理结果TRT 混合推理结果对比项主要包括pred_kp_ofpred_rgbd_segpred_ctr_of这样可以检查ONNX 子图导出后是否出现明显数值漂移TensorRT engine 接入后是否出现异常偏差输入裁剪、点数重采样、choose重映射是否破坏了对齐关系6.3 为什么当前还不能完全脱离tar这也是为什么当前 ORT/TRT 测试还不能算“完全脱离 PyTorch 的纯部署态”。当前tar权重仍然负责支撑未转换的PointNet2 Native CUDA提供完整 PyTorch 基线做精度对照只有当 PointNet2 也被替换为部署后端或者不再需要full_model做对照时tar在测试链里的作用才会进一步缩小。7. ORT 测试脚本解析文件run_linemod_hybrid_ort.py这个脚本的职责不是单独测 ONNX 文件能不能加载而是验证两个 ONNX 子图一段原生 CUDA一套 LINEMOD pose solver能不能被重新拼成一条完整链路。6.1 输入参数的意义关键参数如下--cls指定 LINEMOD 目标物体例如ape--checkpoint用于构建pointnet2_native和full_model--rgb-onnx指向rgb_backbone.onnx--fusion-onnx指向fusion_head.onnx--num-points必须和导出时保持一致当前是4096--height当前固定为480--width当前固定为624--crop-left当前通常是8用于把原始640宽图裁到6246.2 为什么脚本里要先解析绝对路径脚本内部会os.chdir(PVN3D_ROOT)这样做是为了兼容仓库里仍然依赖相对路径的 LINEMOD 旧工具代码。如果不先把以下参数解析成仓库根目录下的绝对路径--checkpoint--rgb-onnx--fusion-onnx--output那么用户从仓库根目录传入的相对路径在chdir之后就会被带偏。这就是之前相对路径找不到文件的原因。6.3 预处理阶段核心函数crop_and_resample_linemod_sample它做了两件必须的事把 RGB 从640裁到624把点数从原始数量重采样到4096与此同时它还必须同步处理chooselabelskp_targ_ofstctr_targ_ofst这里尤其重要的是choose。choose记录的是点与图像像素的对应关系。一旦 RGB 被裁剪像素索引坐标系就变了所以choose必须重映射到新的624宽图坐标上。这也是crop_left参数存在的原因。6.4 ONNX Runtime 部分脚本通过make_sessionrun_rgb_backbonerun_fusion_head分别执行两个 ONNX 子图。具体顺序是rgb - rgb_backbone.onnx - out_rgb, rgb_segcld_rgb_nrm - pointnet2_native - pcld_embout_rgb choose pcld_emb - fusion_head.onnx - pred_*其中 ONNX Runtime 输入输出主要走numpy。6.5 为什么还要构建full_modelORT 脚本不仅构建pointnet2_native还会构建full_model。目的不是重复推理而是为了可选的 PyTorch 数值对比。这个对比用于判断ONNX 导出后数值是否偏离过大裁剪与重采样是否破坏了输入对齐choose重映射是否正确8. TRT 测试脚本解析文件run_linemod_hybrid_trt.pyTRT 脚本和 ORT 脚本在整体流程上是一致的差异只在中间两段执行后端。7.1 输入参数的意义关键参数和 ORT 版基本一致只是--rgb-engine替换了--rgb-onnx--fusion-engine替换了--fusion-onnx其余像--cls、--checkpoint、--num-points、--crop-left的语义不变。7.2TrtEngineRunner的职责TRT 脚本的核心类是TrtEngineRunner它负责读取.engine反序列化 engine创建 execution context根据 binding 信息准备输入输出 tensor调execute_async_v27.3 为什么脚本不依赖pycuda当前实现直接使用torch.cudatorch.Tensor.data_ptr()把 GPU 内存地址传给 TensorRT。也就是说输入 tensor 由torch分配输出 tensor 也由torch分配TensorRT 只消费和写入这些指针这样做的好处是不额外引入pycuda不额外引入cupy可以直接复用当前容器里已经工作的 PyTorch CUDA 环境7.4 TRT 脚本为什么也要构建pointnet2_native和full_model原因和 ORT 版一致pointnet2_native负责补上当前不能转 TensorRT 的那一段full_model负责做可选的 PyTorch 数值对比因此 TRT 脚本不是“纯 TensorRT 整网推理”而是“TensorRT Native CUDA”的混合推理。9. ORT 和 TRT 两个脚本的共同约束8.1 都依赖同一组部署边界无论是 ORT 版还是 TRT 版都必须接受同一组中间张量约定rgb_backbone输出out_rgbpointnet2_native输出pcld_embfusion_head输入out_rgb, choose, pcld_emb只要这个张量契约变了两个测试脚本都要一起改。8.2 都依赖同一组输入尺寸当前已验证的输入配置是height480width624num_points4096crop_left8这些值不是任意选择的而是已经和当前导出产物绑定。8.3 都依赖同一类 checkpoint / 类别语义对 LINEMOD 单物体流程来说num_classes通常固定为2但这不表示导出产物能跨物体复用。例如从ape换到can时通常仍然是--num-classes2但下面这些都要换checkpointONNX 目录engine 目录--cls10. 推荐测试顺序当前不建议一上来就只看 TRT。推荐顺序是先确认model_wrappers.py可以正确加载 checkpoint 并完成拆分先导出 ONNX先跑 ORT 混合测试再构建 TensorRT engine再跑 TRT 混合测试原因很直接ORT 更容易定位图结构与输入输出问题TRT 更适合做最终部署验证先 ORT 后 TRT排错成本最低11. Graphviz 图10.2 图的阅读方式这张图包含两层信息构建产物关系运行时数据流关系构建产物关系说明checkpoint 先进入model_wrappers.pywrapper 拆分后导出 ONNXONNX 再转 TensorRT engineORT/TRT 测试脚本分别消费这些产物运行时数据流说明rgb进入rgb_backbonecld_rgb_nrm进入PointNet2Nativeout_rgb choose pcld_emb进入fusion_headpred_*再进入 pose solver12. 当前结论当前三份代码的职责边界已经比较清晰model_wrappers.py负责定义部署拆分边界并提供统一的子模块构造方式run_linemod_hybrid_ort.py负责验证 ONNX Runtime 版混合链run_linemod_hybrid_trt.py负责验证 TensorRT engine 版混合链它们共同构成了当前 PVN3D 的部署测试主线先拆模型再导出 ONNX先用 ORT 校验子图与联调逻辑再转 TensorRT engine最后用 TRT 验证更接近部署形态的混合推理链