1. 目标本文记录如何让deploy/models/onnx_ape/pvn3d_full.onnx在onnxruntime中真正执行起来并与 PyTorch 原生权重做精度对比。本次不是停留在“脚本骨架”或“待补.so”层面而是在pvn3d-dev容器内实际完成了ORT custom ops 实现full ONNX 重新导出ORT 全图推理PyTorch 精度对比pose / ADD / ADD-S 验证2. 最终采用的实现路线一开始有两条可能路线原生 C / CUDA ORT custom ops.soonnxruntime-extensionsPython custom ops本次最终采用的是onnxruntime-extensionsPython custom ops原因不是偏好问题而是容器里的现状决定了这条路更快能落地。3. 关键上下文梳理3.1 当前 PointNet2 的实际算子当前完整图里要接起来的 PointNet2 custom node 包括PVN3D_FurthestPointSamplePVN3D_GatherPointsPVN3D_BallQueryPVN3D_GroupPointsPVN3D_ThreeNNPVN3D_ThreeInterpolate虽然用户最先强调的是PVN3D_ThreeNNPVN3D_BallQueryPVN3D_FurthestPointSample但为了让整张pvn3d_full.onnx真正执行起来还必须把另外三个也一起补齐否则 ORT 仍然无法跑完整图。3.2 现有 CUDA 语义来源本次 Python custom op 的行为不是拍脑袋写的而是对齐当前仓库已有 CUDA 实现pvn3d/_ext-src/src/sampling_gpu.cupvn3d/_ext-src/src/ball_query_gpu.cupvn3d/_ext-src/src/group_points_gpu.cupvn3d/_ext-src/src/interpolate_gpu.cu因此以下行为都保持和原 CUDA kernel 一致furthest_point_sample首个采样点固定从0开始跳过近零点ball_query邻域不足时重复首个命中索引three_nn先输出dist2图里再接Sqrt4. 容器里遇到的关键问题和解决方案4.1 问题一宿主机没有导出运行依赖在当前宿主机工作区里直接检查时发现torch不存在onnx不存在所以完整导出和 ORT 运行都不能在宿主机侧完成。解决方案统一切到pvn3d-dev容器使用容器内的conda activate pvn3d这一步不是优化而是前提。4.2 问题二完整 ONNX 导出时被动态max_pool2d卡住首次导完整图时报错并不在 PointNet2 custom symbolic 本身而是在F.max_pool2d(new_features, kernel_size[1, new_features.size(3)])这里kernel_size含有动态图尺寸ONNX 导出器要求常量。解决方案把这段改成torch.max(new_features,dim3,keepdimTrue)[0]对应修改文件pvn3d/lib/pointnet2_utils/pointnet2_modules.py这样语义不变但导出友好。4.3 问题三原计划的原生 ORT C custom op 路线被头文件阻塞在容器内检查onnxruntime1.16.3安装内容时发现运行时capi目录存在但开发头文件如onnxruntime_c_api.h、onnxruntime_cxx_api.h不在 wheel 中这意味着不能直接在当前容器里无缝编译原生 ORT custom op.so解决方案切换到onnxruntime-extensions0.10.0用 Python custom op 先把运行链打通。4.3.1 官方 C API 路线应该怎么做虽然这次先用onnxruntime-extensions跑通了链路但如果后续要回到 ONNX Runtime 官方 C API 路线建议按下面的方式做。A. 官方 C API 依赖什么官方 C custom op 路线依赖的是 ONNX Runtime 提供的 C / C 头文件和运行时库核心通常包括onnxruntime_c_api.honnxruntime_cxx_api.hOrt::CustomOpDomainOrt::SessionOptions对应官方文档入口可以参考ONNX Runtime C APIOrt::CustomOpDomainhttps://onnxruntime.ai/docs/api/c/struct_ort_1_1_custom_op_domain.html这条路线的核心思路是用 C 实现OrtCustomOp或其封装在自定义库里注册 custom op通过SessionOptions.RegisterCustomOpsLibrary(...)或手动CustomOpDomain.Add(...)让 ORT runtime 在执行图时调到你的原生实现B. 当前环境下必须匹配的版本如果要走官方 C API 路线最重要的原则不是“越新越好”而是头文件版本和运行时版本必须严格对齐结合你当前容器建议锁定为ONNX Runtime runtime1.16.3ONNX Runtime C / C headers1.16.3Python3.8CUDA沿用当前容器里的可用版本链不要在当前容器里出现这种错配runtime 是1.16.3头文件拿1.18.x或main分支这样非常容易在API 结构体布局custom op 注册接口ABI上出问题。C. 当前容器里缺的不是 runtime而是 dev headers当前pvn3d-dev容器内已经有onnxruntime1.16.3但缺少onnxruntime_c_api.honnxruntime_cxx_api.h其他开发用头文件也就是说当前容器里的 pip wheel 更像是运行时安装不是完整的 C 开发包。D. 官方 C API 路线在当前容器里的推荐操作方法如果你要继续推进原生 C ORT custom op建议按下面顺序操作。步骤 1拿到完全匹配的 ORT 1.16.3 开发头文件建议使用以下任一方式下载 ONNX Runtime1.16.3官方 release 包并提取include/下载 ONNX Runtimev1.16.3源码 tag并使用其中的公开头文件这里最关键的要求只有一个必须是1.16.3不要混用其他版本。步骤 2确认运行时库来源当前容器里已经有 ORT runtime主要在 Python 包目录下。你至少要确认onnxruntime当前 Python 运行时能正常 importSessionOptions.register_custom_ops_library(...)可用这两点当前已经在容器内确认过。步骤 3新建原生 ORT custom op 工程目录建议单独建一套目录例如deploy/ort_custom_ops/ CMakeLists.txt include/ src/ three_nn_op.cc ball_query_op.cc furthest_point_sample_op.cc register.cc不要把 ORT custom op 和 TensorRT plugin 混在一个目录里。原因很简单二者运行时不同注册接口不同调试方式不同步骤 4优先实现最小 3 个算子官方 C API 路线下仍然建议优先顺序PVN3D_ThreeNNPVN3D_BallQueryPVN3D_FurthestPointSample原因和之前一样ThreeNN边界最清楚BallQuery可以尽早固定邻域补位语义FurthestPointSample是 SA 路径最核心的离散算子步骤 5注册 custom domain原生 C 路线建议在注册层显式创建Ort::CustomOpDomain然后把每个OrtCustomOp加进去。如果继续沿用当前已经跑通的 full ONNX建议保持 domain 为ai.onnx.contrib这样就不需要再次改图。如果后面不再依赖onnxruntime-extensions而是完全切换到原生 C.so也可以重新评估是否把导出 domain 改回私有域名。步骤 6通过register_custom_ops_library装载最终运行脚本侧建议继续保持SessionOptions.register_custom_ops_library(...)这样 Python 入口层不需要大改只要把当前 Python custom op 路线替换成原生.so就行。E. 当前最推荐的过渡策略对当前项目来说最稳的路线不是立刻把 Python custom op 全部推倒重写而是保留当前onnxruntime-extensionsPython custom op 作为语义基线用它做功能验证和精度对齐再逐个把性能敏感算子替换成官方 C API custom op这样做的好处是你已经有一条真实可运行的 ORT 链路原生 C 路线出问题时有现成基线可对照每个算子的语义偏差更容易定位4.4 问题四onnxruntime-extensions只能稳定承接默认 custom domain仓库最开始给 full ONNX 导出设置的 custom domain 是com.pvn3d.pointnet2但容器里检查onnxruntime-extensions后发现它的默认 domain 是ai.onnx.contrib并且 Python custom op 注册链是围绕这个 domain 工作的。解决方案把导出侧 custom domain 改成ai.onnx.contrib对应文件deploy/scripts/pointnet2_custom_ops.py然后重新导出 full ONNX。4.5 问题五最初的 ORT 运行脚本设计依赖外部.so早期脚本是按SessionOptions.register_custom_ops_library(/path/to/libxxx.so)这种“外部 ORT custom op 动态库”方案写的。但切换到onnxruntime-extensions以后这个假设就不再适合当前阶段。解决方案改成在 Python 侧用onnx_op注册 PointNet2 custom opimport这些 Python custom op 模块用onnxruntime_extensions.get_library_path()注册 extensions 自带动态库这样就能让 ORT 通过 Python custom op 桥接运行这些节点。5. 这次实际补齐的代码5.1 custom op symbolic / 导出侧涉及文件deploy/scripts/pointnet2_custom_ops.pydeploy/scripts/export_full_onnx.pydeploy/model_wrappers.py主要完成PointNet2 custom node symbolic 注册full ONNX 导出onnx.checkeronnx.shape_inference节点清单和报告输出5.2 ORT Python custom op 实现新增文件deploy/scripts/ort_pointnet2_pyops.py实际实现了 6 个 ORT Python custom opPVN3D_FurthestPointSamplePVN3D_GatherPointsPVN3D_BallQueryPVN3D_GroupPointsPVN3D_ThreeNNPVN3D_ThreeInterpolate其中前三个是最初优先关注的核心算子后三个是为了让 full graph 真正跑起来一并补齐的必要项。5.3 ORT 全图运行脚本涉及文件deploy/scripts/run_full_onnx_ort.py脚本能力包括加载pvn3d_full.onnx注册onnxruntime-extensionslibrary导入 Python custom ops跑完整 ORT 前向与 PyTorch 原生模型对比可选继续做 pose / ADD / ADD-S6. 实际运行环境本次真实运行环境容器pvn3d-devcondapvn3dPyTorch1.10.0cu113ONNX1.14.1ONNX Runtime1.16.3onnxruntime-extensions0.10.0本次运行图deploy/models/onnx_ape/pvn3d_full.onnx使用配置clsapesample_index0num_points4096height480width624crop_left87. 实际运行结果本次在容器内已经实际执行成功ORT full graph 前向成功PyTorch 输出对比成功pose / ADD / ADD-S成功实际输出摘要文件deploy/benchmarks/linemod_ape_full_onnx_ort.json本次实际精度对比结果为pred_kp_ofmax_abs 0.008136272430419922mean_abs 7.063127213768894e-06pred_rgbd_segmax_abs 0.031524658203125mean_abs 0.0004335300764068961pred_ctr_ofmax_abs 0.001704871654510498mean_abs 4.566256848193007e-06本次实际 pose 指标ADD 0.0044791861437261105ADD-S 0.0028379589784890413这说明当前 Python custom op 路线已经足够把full ONNX 图ORT runtimePyTorch 精度对比整条链先跑通。8. 当前实现边界虽然这次已经跑通但边界也要写清楚。8.1 当前是 Python custom op不是高性能 C/CUDA ORT kernel当前 ORT full graph 是能跑通的但执行后端是CPUExecutionProviderPython custom op所以这条路线当前更适合作为功能验证图正确性验证custom op 语义对齐验证不适合作为最终性能后端。8.2 当前 domain 已切换为ai.onnx.contrib这是为了和onnxruntime-extensions的 Python custom op 运行机制直接兼容。这条选择是当前环境下的务实方案不是唯一理论方案。8.3 当前 full ORT 路线已完成但 TRT plugin 路线仍然独立这次完成的是ORT full graph 调试链它不等价于TensorRT plugin 已完成TRT full engine 已完成后面如果继续做 TRT仍然要单独推进 plugin 和 engine。9. 最终结论这次工作已经把 ORT custom ops 这条链从“设计阶段”推进到了“真实可运行阶段”ThreeNNBallQueryFurthestPointSample这三个核心算子已经连同GatherPointsGroupPointsThreeInterpolate一起完成了 ORT Python custom op 版本因此pvn3d_full.onnx已经能在 ORT 中真实运行也已经能和 PyTorch 原生权重做输出精度对比后续如果继续推进下一步最合理的是保留当前 Python custom op 作为语义基线再把性能敏感的几个算子逐步下沉成原生 C / CUDA ORT custom op