PVN3D ORT CUDA Custom Ops 实现与联调记录

1. 任务目标本次任务是在已经可用的 CPU C custom op 基础上单独落一条 GPU 版本新建独立目录deploy/ort_custom_ops_gpu/把 PointNet2 的 6 个 custom op 升级为 CUDA custom kernel通过 ONNX Runtime 官方 C custom op API 导出.so在 Python 侧通过register_custom_ops_library(...)加载这份 GPU.so用CUDAExecutionProvider跑通完整 pvn3d_full.onnx与 PyTorch 原生权重做输出和 pose 精度对比这次工作不是替换原来的 CPU 版本而是保持两套并行工程CPU 版deploy/ort_custom_ops/GPU 版deploy/ort_custom_ops_gpu/2. 最终产物2.1 GPU custom op 工程CMakeLists.txtcuda_utils.hpvn3d_pointnet2_kernels.hpvn3d_pointnet2_kernels.cupvn3d_pointnet2_ops_cuda.cc2.2 GPU 联调脚本run_full_onnx_ort_cpp_gpu.py2.3 构建结果libpvn3d_ort_custom_ops_gpu.so2.4 运行结果linemod_ape_full_onnx_ort_cpp_gpu.json3. 环境事实本次以pvn3d-dev容器为准。关键环境Ubuntu 18.04.5g 7.5.0nvcc 11.3.109GPU:NVIDIA GeForce RTX 4060 Laptop GPUcuDNN 8.2onnxruntime1.16.3onnxruntime-gpu1.16.3onnx1.14.1torch1.10.0cu113本次 ORT 开发头文件仍然来自你已经解压好的官方包/workspace/tmp/onnxruntime-linux-x64-1.16.3/include/onnxruntime_c_api.h/workspace/tmp/onnxruntime-linux-x64-1.16.3/include/onnxruntime_cxx_api.h也就是说编译用头文件onnxruntime-linux-x64-1.16.3Python 运行时onnxruntime-gpu1.16.34. 为什么需要单独的 GPU 版本目录原来的 CPU custom op 工程已经能跑完整图但它的实现方式是custom op kernel 在 CPU 上执行ORT provider 走CPUExecutionProvider这次要做的是custom op kernel 在 CUDA 上执行ORT provider 显式走CUDAExecutionProvider因此必须单独拆目录避免以下混淆CMake 构建语言不同运行时 provider 绑定不同.so依赖不同调试问题完全不同5. GPU 版实现策略5.1 沿用 ORT 官方CustomOpBase本次没有切换到另一套封装而是继续沿用Ort::CustomOpBaseTOp, TKernel区别在于每个 op 都显式声明constchar*GetExecutionProviderType()const{returnCUDAExecutionProvider;}这样 ORT 才会把这些 custom nodes 分配到 CUDA EP。5.2 从 ORT 取当前 CUDA stream每个 kernel 运行时都通过Ort::KernelContextctx(context);auto*streamstatic_castcudaStream_t(ctx.GetGPUComputeStream());拿到 ORT 当前执行流。这是这次 GPU custom op 能和 ORT CUDA EP 正常配合的关键点。5.3 直接复用 PVN3D PointNet2 的 CUDA 语义GPU kernel 不是重新设计的而是直接对齐当前仓库已有 PointNet2 CUDA 实现。对应来源主要是sampling_gpu.cuball_query_gpu.cugroup_points_gpu.cuinterpolate_gpu.cu但本次做了一个重要适配原仓库很多 kernel 用的是intONNX / ORT 这边图里索引统一是int64所以 GPU 版.so里把索引相关输入输出都统一改成了int64_t。6. 本次承接的 custom opGPU 版一次性承接了完整图里全部 6 个 PointNet2 custom opPVN3D_FurthestPointSamplePVN3D_GatherPointsPVN3D_BallQueryPVN3D_GroupPointsPVN3D_ThreeNNPVN3D_ThreeInterpolate所有 op 都注册在custom domain:ai.onnx.contrib这和当前 full ONNX 导出结果保持一致。7. 容器内操作过程7.1 检查 CUDA 与 ORT 初始状态最开始容器里虽然有 CUDA 工具链和 GPU但 Python 环境里的 ORT 只有 CPU provider。实际检查命令dockerexecpvn3d-devbash-lc which nvcc || true nvcc --version || true source /opt/conda/etc/profile.d/conda.sh conda activate pvn3d python - PY import onnxruntime as ort print(ort, ort.__version__) print(providers, ort.get_available_providers()) PY 当时结果是nvcc 11.3.109ort 1.16.3providers 只有CPUExecutionProvider7.2 安装onnxruntime-gpu为了让 Python 入口层能够真正使用 CUDA EP实际在容器内安装了dockerexecpvn3d-devbash-lc source /opt/conda/etc/profile.d/conda.sh conda activate pvn3d python -m pip install --upgrade --force-reinstall onnxruntime-gpu1.16.3 安装后再次检查dockerexecpvn3d-devbash-lc source /opt/conda/etc/profile.d/conda.sh conda activate pvn3d python - PY import onnxruntime as ort print(ort, ort.__version__) print(providers, ort.get_available_providers()) PY 结果变为TensorrtExecutionProviderCUDAExecutionProviderAzureExecutionProviderCPUExecutionProvider7.3 编译 GPU custom op.so最终使用的构建命令dockerexecpvn3d-devbash-lc cd /workspace/workflow/self/PVN3D/deploy/ort_custom_ops_gpu rm -rf build mkdir build cd build cmake -DONNXRUNTIME_ROOT/workspace/tmp/onnxruntime-linux-x64-1.16.3 .. cmake --build . -- -j2 构建成功后得到libpvn3d_ort_custom_ops_gpu.so7.4 检查导出符号和链接检查命令dockerexecpvn3d-devbash-lc cd /workspace/workflow/self/PVN3D nm -D deploy/ort_custom_ops_gpu/build/libpvn3d_ort_custom_ops_gpu.so | rg RegisterCustomOps -n -S ldd deploy/ort_custom_ops_gpu/build/libpvn3d_ort_custom_ops_gpu.so 结果说明RegisterCustomOps已正确导出.so已正确依赖libcudart.so.11.07.5 先做 smoke test先跳过 PyTorch 对比和 pose只验证.so能否被 ORT 加载ORT 是否真的能用 CUDA EP 跑完整图命令dockerexecpvn3d-devbash-lc cd /workspace/workflow/self/PVN3D source /opt/conda/etc/profile.d/conda.sh conda activate pvn3d python deploy/scripts/run_full_onnx_ort_cpp_gpu.py \ --checkpoint weights/ape_pvn3d_best.pth.tar \ --onnx deploy/models/onnx_ape/pvn3d_full.onnx \ --custom-ops-lib deploy/ort_custom_ops_gpu/build/libpvn3d_ort_custom_ops_gpu.so \ --cls ape \ --sample-index 0 \ --num-points 4096 \ --height 480 \ --width 624 \ --crop-left 8 \ --skip-torch-compare \ --skip-pose-eval 这一步成功说明CUDAExecutionProvider可用.so可被 ORT 正常加载6 个 custom nodes 可被 CUDA EP 承接完整pvn3d_full.onnx已可执行7.6 再做完整验证命令dockerexecpvn3d-devbash-lc cd /workspace/workflow/self/PVN3D source /opt/conda/etc/profile.d/conda.sh conda activate pvn3d python deploy/scripts/run_full_onnx_ort_cpp_gpu.py \ --checkpoint weights/ape_pvn3d_best.pth.tar \ --onnx deploy/models/onnx_ape/pvn3d_full.onnx \ --custom-ops-lib deploy/ort_custom_ops_gpu/build/libpvn3d_ort_custom_ops_gpu.so \ --cls ape \ --sample-index 0 \ --num-points 4096 \ --height 480 \ --width 624 \ --crop-left 8 \ --output deploy/benchmarks/linemod_ape_full_onnx_ort_cpp_gpu.json 8. 实际结果最终结果写入linemod_ape_full_onnx_ort_cpp_gpu.json关键结果pred_kp_ofmax_abs 0.014214515686035156mean_abs 8.158481250575278e-06pred_rgbd_segmax_abs 0.07232666015625mean_abs 0.0005051378393545747pred_ctr_ofmax_abs 0.001059534028172493mean_abs 5.723987214878434e-06ADD 0.004475891590118408ADD-S 0.002825918374583125当前这组结果说明GPU custom op 路线已可执行完整图输出与 PyTorch 仍保持较小偏差pose 层面的 ADD / ADD-S 结果也稳定9. 中间遇到的问题与处理9.1 问题一容器里虽然有 CUDA但 ORT 只有 CPU provider现象ort.get_available_providers()只有CPUExecutionProvider结论有nvcc和 GPU 不代表 Python 侧 ORT 就能跑 CUDA处理安装onnxruntime-gpu1.16.39.2 问题二GPU custom op 的.cc文件也需要 CUDA 头文件现象第一次编译时.cu可以编但.cc编译报错fatal error: cuda_runtime.h: No such file or directory原因pvn3d_pointnet2_ops_cuda.cc里同样会拿cudaStream_t但 CMake 只让 CUDA 编译单元看到了 CUDA include path处理在target_include_directories(...)中显式补/usr/local/cuda/include9.3 问题三CMake 3.10 对 CUDA arch 参数传递不干净现象第一版虽然能编但 device link 阶段出现nvlink warning : SM Arch (sm_52) not found ...原因容器里是cmake 3.10.2仅靠 target 级别的-gencode参数device link 阶段仍可能带上默认旧架构处理把-gencode放到CMAKE_CUDA_FLAGS重新构建后告警消失当前固定配置为sm_86compute_86这样在当前nvcc 11.3下对 RTX 4060 这类新卡更稳。9.4 问题四原始 PVN3D CUDA kernel 使用int但 ONNX 图里是int64现象原 PointNet2 CUDA 扩展里的索引多为int但当前 ONNX symbolic 导出的索引节点类型是int64如果直接照搬原 kernel会在 ORT custom op 输入输出类型上不一致。处理GPU custom op 版本统一把索引相关 tensor 改成int64_t涉及FurthestPointSample输出BallQuery输出ThreeNN的idx输出GatherPoints / GroupPoints / ThreeInterpolate的idx输入9.5 问题五GPU custom op 不能继续沿用 CPU 版运行脚本现象CPU 版脚本默认 provider 是CPUExecutionProvider也不会检查CUDAExecutionProvider是否真的可用处理单独新建 run_full_onnx_ort_cpp_gpu.py默认 provider 改为[CUDAExecutionProvider, CPUExecutionProvider]并且显式校验CUDAExecutionProvider in ort.get_available_providers()10. 当前边界这次已经完成的是独立 GPU custom op 工程6 个 PointNet2 custom op 的 CUDA kernel 版本原生.so构建Python 侧加载 GPU.soCUDAExecutionProvider下完整图执行与 PyTorch 的输出与 pose 对比这次还没有做的是CUDA kernel 的性能优化多 batch / dynamic shape 调优ORT 自定义 allocator / workspace 池化Nsight 级别 profiling当前实现定位是正确性优先的 GPU custom op 基线为后续性能优化和 TensorRT plugin 对齐提供参考11. 结论这次工作已经把 ORT custom op 这条链从CPU C custom op 可运行推进到了CUDA C custom op 可运行能在 ORTCUDAExecutionProvider下执行完整pvn3d_full.onnx能与 PyTorch 原生结果做稳定对比后续如果继续推进优先顺序建议是对 6 个 kernel 做 profiling识别最耗时的 PointNet2 op再决定是否继续优化 ORT GPU 版或回到 TensorRT plugin 路线