当前位置：首页 > news >正文

PVN3D Custom ONNX Op / TensorRT Plugin 设计说明

news 2026/7/14 17:00:32

1. 目标

本文用于给 PVN3D 的PointNet2自定义 CUDA 算子定义一套可落地的部署侧设计，目标是：

支持导出完整ONNX
支持构建完整TensorRT engine
部署运行时不再依赖PyTorch和pvn3d.lib.pointnet2_utils._ext

本文不是泛泛而谈的 PointNet2 设计文档，而是直接结合当前仓库里的真实实现：

Python 封装：
- pointnet2_utils.py
- pointnet2_modules.py
C++/CUDA 扩展：
- bindings.cpp
- sampling.cpp
- sampling_gpu.cu
- ball_query.cpp
- ball_query_gpu.cu
- group_points.cpp
- group_points_gpu.cu
- interpolate.cpp
- interpolate_gpu.cu

2. 版本基线

本文对应的部署基线沿用当前容器文档中的已核验版本：

Python3.8.20
PyTorch1.10.0+cu113
ONNX1.14.1
ONNX Runtime1.16.3
TensorRT8.6.1

这里的ONNX -> TRT默认指：

TensorRT8.6.1自带 ONNX parser
trtexec或 TensorRT builder API

而不是依赖单独的onnx-tensorrtPython 包。

3. 现有 CUDA 扩展暴露出的算子集合

从 bindings.cpp 可以确认，当前_ext实际导出的是：

gather_points
gather_points_grad
furthest_point_sampling
three_nn
three_interpolate
three_interpolate_grad
ball_query
group_points
group_points_grad

对部署前向链真正有意义的是：

furthest_point_sampling
gather_points
ball_query
group_points
three_nn
three_interpolate

*_grad相关接口只服务训练反向，对部署态完整 ONNX / TRT engine 不必实现。

4. 场景描述：这些算子在 PVN3D 里到底干什么

4.1 场景 A：Set Abstraction 路径

对应模块：

pointnet2_modules.py

在PointnetSAModuleMSG.forward()里，PointNet2 做的是：

从原始点云xyz[B,N,3]中选出npoint个中心点
以这些中心点为球心，在半径radius内搜邻域点
对每个中心点收集固定数量nsample的邻域索引
用这些索引把xyz和features分组为局部 patch
对局部 patch 做 MLP 和池化

也就是说，这条链不是普通 CNN 的卷积 receptive field，而是点云上的“显式邻域构造”：

furthest_point_sample
gather_operation
ball_query
grouping_operation

这是第一类场景。

4.2 场景 B：Feature Propagation 路径

对应模块：

pointnet2_modules.py

在PointnetFPModule.forward()里，PointNet2 做的是：

对稀疏点集和稠密点集做 3 近邻搜索
计算三近邻的距离倒数权重
按三近邻索引把稀疏特征插值回稠密点集

对应链路是：

three_nn
three_interpolate

这是第二类场景。

4.3 为什么这两类场景必须单独描述

因为它们决定了 custom op 的接口边界：

SA路径主要是“离散索引生成 + batched gather”
FP路径主要是“近邻搜索 + 加权插值”

如果不先按场景分开，后面容易把 plugin 设计做成一堆彼此耦合、难以调试的杂合层。

5. 从 CUDA 源码提炼出来的真实语义

这部分非常关键，因为 custom op 的行为必须和现有 CUDA kernel 保持一致。

5.1`furthest_point_sampling`

从 sampling_gpu.cu 可以确认：

输入是dataset[b, n, 3]
输出是idxs[b, m]
首个采样点固定从old = 0开始
temp[b, n]初始化为大数1e10
迭代过程中使用“当前点到已选点集合的最小距离”的最大值作为下一个点
对模长mag <= 1e-3的点直接跳过

这意味着第一版 custom op 设计里要明确：

采样起点不是随机的，而是固定0
存在“近零点跳过”行为
算子行为依赖固定的距离更新逻辑，而不是任意 FPS 变体

5.2`ball_query`

从 ball_query_gpu.cu 可以确认：

输入是：
- new_xyz[b, m, 3]
- xyz[b, n, 3]
输出是idx[b, m, nsample]
判定条件是d2 < radius^2
当找到第一个命中点k时，会先把整个idx[j, :]全部填成k
后续继续找到命中点时，再按cnt位置覆写前面的槽位

这意味着：

邻域不足nsample时，不是填-1
而是重复第一个命中的索引

这是非常重要的行为约定，custom op 和 TRT plugin 必须保持一致。

5.3`group_points`

从 group_points_gpu.cu 可以确认：

输入：
- points[b, c, n]
- idx[b, npoints, nsample]
输出：
- out[b, c, npoints, nsample]

本质就是按idx从points上取值并重排。

它本身没有搜索语义，依赖上游idx的定义是否一致。

5.4`gather_points`

从 sampling_gpu.cu 可以确认：

输入：
- points[b, c, n]
- idx[b, m]
输出：
- out[b, c, m]

它对应的是单索引 gather，而不是group_points那种二维索引 gather。

5.5`three_nn`

从 interpolate_gpu.cu 可以确认：

输入：
- unknown[b, n, 3]
- known[b, m, 3]
输出：
- dist2[b, n, 3]
- idx[b, n, 3]
kernel 内输出的是平方距离dist2
Python 封装层再对dist2做torch.sqrt()

这意味着部署设计里有两个选择：

custom op 直接输出sqrt(dist2)，对齐 Python 最终语义
custom op 输出dist2，再在图里补一个Sqrt

推荐第一版选择第 2 种：

plugin 更贴近现有 CUDA kernel
方便先保持实现简单

5.6`three_interpolate`

从 interpolate_gpu.cu 可以确认：

输入：
- points[b, c, m]
- idx[b, n, 3]
- weight[b, n, 3]
输出：
- out[b, c, n]

其语义严格就是：

三个索引位置的特征值乘对应权重后求和

这意味着如果上游three_nn已经产出稳定的idx，那么three_interpolate比搜索类算子更规则。

6. Custom ONNX op 域与命名

建议统一使用固定 domain：

com.pvn3d.pointnet2

建议的 op 名如下：

PVN3D_FurthestPointSample
PVN3D_GatherPoints
PVN3D_BallQuery
PVN3D_GroupPoints
PVN3D_ThreeNN
PVN3D_ThreeInterpolate

这样做的目的：

和 TensorRT plugin 名一一对应
后续图检查时容易定位
不和其他第三方 PointNet2 实现混淆

7. 每个 custom op 的接口设计

7.1`PVN3D_FurthestPointSample`

场景：

SA路径选中心点

输入：

xyz:float[B, N, 3]

属性：

npoint: int

输出：

idx:int32[B, npoint]

约束：

第一版只支持 static shape
第一版只支持B=1
第一版只支持float32输入

语义约定：

起始索引固定为0
保持现有 kernel 的“近零点跳过”行为

TensorRT plugin 建议名：

PVN3DFurthestPointSample_TRT

7.2`PVN3D_GatherPoints`

场景：

SA路径根据 FPS 输出索引取中心点坐标

输入：

points:float[B, C, N]
idx:int32[B, S]

输出：

out:float[B, C, S]

第一版建议：

先尝试标准 ONNXGather改写
只有在 TRT parser 或运行期不稳定时才升级成 plugin

如果最终也做 plugin，plugin 名建议：

PVN3DGatherPoints_TRT

7.3`PVN3D_BallQuery`

场景：

SA路径在每个中心点周围搜固定邻域

输入：

new_xyz:float[B, S, 3]
xyz:float[B, N, 3]

属性：

radius: float
nsample: int

输出：

idx:int32[B, S, nsample]

语义约定：

距离判定条件是d2 < radius^2
邻域不足时重复第一个命中索引
不使用-1作为 padding

TensorRT plugin 建议名：

PVN3DBallQuery_TRT

7.4`PVN3D_GroupPoints`

场景：

SA路径把邻域索引转换成局部 patch 特征张量

输入：

points:float[B, C, N]
idx:int32[B, S, K]

输出：

out:float[B, C, S, K]

第一版建议：

先尝试标准图改写
如果失败，再做 plugin

如果最终 plugin 化，plugin 名建议：

PVN3DGroupPoints_TRT

7.5`PVN3D_ThreeNN`

场景：

FP路径从稀疏点集向稠密点集回传特征前，先找 3 近邻

输入：

unknown:float[B, N, 3]
known:float[B, M, 3]

输出：

dist2:float[B, N, 3]
idx:int32[B, N, 3]

说明：

推荐 custom op 直接输出dist2
后续图里补Sqrt
这样最贴近当前 CUDA kernel

TensorRT plugin 建议名：

PVN3DThreeNN_TRT

7.6`PVN3D_ThreeInterpolate`

场景：

FP路径根据三近邻和权重做特征插值

输入：

points:float[B, C, M]
idx:int32[B, N, 3]
weight:float[B, N, 3]

输出：

out:float[B, C, N]

第一版建议：

优先尝试标准图改写
若在 TRT 8.6.1 下不稳定，再实现 plugin

如果最终 plugin 化，plugin 名建议：

PVN3DThreeInterpolate_TRT

8. 第一版实现边界

为了降低复杂度，第一版 custom op / plugin 明确建议限制为：

static shape
batch=1
num_points=4096
输入特征只支持float32
索引统一int32
先不考虑反向

不要第一版就追求：

dynamic shape
batch>1
fp16kernel 内部全支持
完整训练反向

这些优化都应该放在完整链路先打通之后。

9. ONNX symbolic 设计建议

PyTorch 侧建议为每个算子注册独立 symbolic，直接映射到 custom domain 节点。

例如：

furthest_point_sample(xyz, npoint)
->com.pvn3d.pointnet2::PVN3D_FurthestPointSample
ball_query(radius, nsample, xyz, new_xyz)
->com.pvn3d.pointnet2::PVN3D_BallQuery
three_nn(unknown, known)
->com.pvn3d.pointnet2::PVN3D_ThreeNN

原则：