cuda-pointpillar+openpcdet 0.6.0 自定义模型适配指北
目录
一、整体架构概览
整个推理 pipeline 分为 4 个阶段,但 ONNX 只导出中间 2 个:
┌──────────────┐ ┌──────────────────┐ ┌──────────────────┐ ┌──────────────┐
│ Voxelization │ → │ PillarEncoder │ → │ 2D Backbone │ → │ PostProcess │
│ (预处理,CPU) │ │ + ScatterBEV │ │ + DetectHead │ │ (NMS,CPU) │
│ │ │ ← ONNX 导出这段 →│ │ ← ONNX 导出这段 →│ │ │
└──────────────┘ └──────────────────┘ └──────────────────┘ └──────────────┘
阶段说明:
-
Voxelization:将原始 3D 点云转换为体素(voxel)表示,由 CUDA/C++ 预处理模块完成,不在 ONNX 中
-
PillarEncoder + ScatterBEV:提取每个柱子(pillar)的特征,散射到 2D 伪图像,ONNX 导出
-
2D Backbone + DetectHead:2D 卷积骨干网络 + 检测头输出分类/框回归/方向预测,ONNX 导出
-
PostProcess:NMS(非极大值抑制)后处理,由 CUDA/C++ 模块完成,不在 ONNX 中
二、export_onnx.py — 从 OpenPCDet 导出 ONNX
Step 1: 构建模型 + 加载权重(第 86-99 行)
model = build_network(model_cfg=cfg.MODEL, num_class=len(cfg.CLASS_NAMES), dataset=demo_dataset)
model.load_params_from_file(filename=args.ckpt, logger=logger, to_cpu=True)
model.cuda().eval()
说明:
-
build_network():根据配置文件(cfgs/kitti_models/pointpillar.yaml)构建 PointPillar 模型架构 -
load_params_from_file():加载训练好的.ckpt权重文件 -
.cuda().eval():将模型移至 GPU 并设置为评估模式(禁用 Dropout、BatchNorm 更新等)
Step 2: 构造 dummy 输入(第 103-126 行)
MAX_VOXELS = 10000
dummy_voxels = torch.zeros((10000, 32, 4), dtype=torch.float32)
dummy_voxel_idxs = torch.zeros((10000, 4), dtype=torch.int32)
dummy_voxel_num = torch.zeros((1,), dtype=torch.int32)
输入张量含义:
| 张量 | 形状 | 数据类型 | 说明 |
|------|------|----------|------|
| voxels | [V, P, C] | float32 | V=最大体素数(10000),P=每体素最大点数(32),C=特征维度(4 = x, y, z, intensity) |
| voxel_idxs | [V, 4] | int32 | 体素坐标 (batch_id, z, y, x) |
| voxel_num | [1] | int32 | 实际有效体素数量 |
为什么要用 dummy 输入:
-
torch.onnx.export需要执行一次前向传播来追踪计算图 -
ONNX 输入的形状固定,所以用最大容量(10000 体素)作为静态形状
-
实际推理时无效体素(
voxel_num之外的)会被 mask 掉
Step 3: 导出原始 ONNX(第 128-137 行)
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 模型输入(用于追踪计算图)
"pointpillar_raw.onnx", # 输出文件路径
export_params=True, # 导出训练权重
opset_version=11, # ONNX 算子集版本
do_constant_folding=True, # 常量折叠优化
keep_initializers_as_inputs=True,
input_names=['voxels', 'voxel_num', 'voxel_idxs'],
output_names=['cls_preds', 'box_preds', 'dir_cls_preds']
)
输出张量含义:
| 张量 | 形状 | 说明 |
|------|------|------|
| cls_preds | [1, 248, 216, 18] | 分类预测:18 = 3类 × 6anchor |
| box_preds | [1, 248, 216, 42] | 框回归:42 = 7 × 6anchor |
| dir_cls_preds | [1, 248, 216, 12] | 方向分类:12 = 2 × 6anchor |
问题:
-
原始 ONNX 包含了 OpenPCDet 完整的计算图,包括很多在 TensorRT 中不兼容/不需要的操作
-
后处理(NMS、坐标解码等)由 C++ 模块完成,不应该在 ONNX 中
-
需要专门的修改和简化
Step 4: 三阶段 ONNX 修改流水线(第 139-146 行)
# 阶段 1:截断后处理
onnx_trim_post = simplify_postprocess(onnx_raw)
# 阶段 2:图简化(常数折叠、死代码消除等)
onnx_simp, check = simplify(onnx_trim_post)
assert check, "Simplified ONNX model could not be validated"
# 阶段 3:预处理重构 + ScatterBEV 替换为 Plugin
onnx_final = simplify_preprocess(onnx_simp)
onnx.save(onnx_final, "pointpillar.onnx")
三、modify_onnx.py — ONNX 图手术详解
函数 1: simplify_postprocess(第 40-74 行)
目的:截断后处理操作,只保留 2D backbone 到检测头输出的部分。
详细步骤:
- 定义最终输出节点(第 44-46 行):
```python
cls_preds = gs.Variable(name="cls_preds", shape=(1, 248, 216, 18))
box_preds = gs.Variable(name="box_preds", shape=(1, 248, 216, 42))
dir_cls_preds = gs.Variable(name="dir_cls_preds", shape=(1, 248, 216, 12))
```
- 定位关键节点(第 59-67 行):
```python
# 找到第一个 ConvTranspose(上采样层的起点)
first_ConvTranspose_node = [node for node in graph.nodes if node.op == "ConvTranspose"][0]
# 往下走 3 步,找到 Concat 节点(三个检测头特征拼接点)
concat_node = loop_node(graph, first_ConvTranspose_node, 3)
# 找到 Concat 后的 3 个子图分支,每条走 1 步到 Transpose
first_node_after_concat = [...]
for i in range(3):
transpose_node = loop_node(graph, first_node_after_concat[i], 1)
transpose_node.outputs = [new_outputs[i]]
```
- 清理输入/输出(第 52-71 行):
- 删除除 voxels/voxel_idxs/voxel_num 外的所有输入
- 将 3 个 Transpose 节点的输出重定义为最终输出
效果对比:
原始 OpenPCDet 图:
... → ConvTranspose → Concat → Transpose → [后处理一堆节点] → cls_preds
└→ Transpose → [后处理一堆节点] → box_preds
└→ Transpose → [后处理一堆节点] → dir_cls_preds
修改后的 ONNX 图:
... → ConvTranspose → Concat → Transpose → cls_preds
└→ Transpose → box_preds
└→ Transpose → dir_cls_preds
函数 2: simplify_preprocess(第 77-142 行)
这是最核心的修改,做了两件重要的事情。
修改 1:扩展输入特征维度(第 88 行)
input_new = gs.Variable(name="voxels", dtype=np.float32, shape=(MAX_VOXELS, 32, 10))
# ↑ 从 4 改为 10
原理:
-
原始 OpenPCDet 输入
voxels的特征维度是 4(x, y, z, intensity) -
VFE(Voxel Feature Encoder)会计算额外的特征(如点到体素中心的偏移量),最终每点 10 维特征
-
ONNX 导出时这 4→10 的扩展在图内,通过修改将输入直接改为 10 维
-
意味着 Voxelization 的特征扩展在导出前由外部完成
修改 2:手动重构 PillarEncoder → Scatter 的计算链(第 99-136 行)
原始 OpenPCDet 计算图:
voxels[V,32,4]
↓ [特征扩展]
voxels[V,32,10]
↓ [MatMul 线性映射]
features[V*32,64]
↓ [BatchNorm + ReLU]
features[V*32,64]
↓ [ReduceMax 沿点维度]
pillar_features[V,64]
↓ [ScatterBEV 操作]
spatial_features[1,64,496,432]
↓ [2D Backbone]
cls/box/dir_preds
修改后的 ONNX 计算图:
voxels[V,32,10] (直接输入)
↓ [Reshape]
features_flat[V*32,10]
↓ [MatMul @ 权重矩阵 10×64]
features[V*32,64]
↓ [BatchNorm]
features[V*32,64]
↓ [ReLU]
features[V*32,64]
↓ [Reshape]
features_reshaped[V,32,64]
↓ [ReduceMax 沿点维度取最大值]
pillar_features[V,64]
↓ [PPScatterPlugin ← 自定义 TensorRT Plugin]
spatial_features[1,64,496,432]
↓ [2D Backbone]
cls/box/dir_preds
关键代码解析:
- Reshape 操作(第 99-105 行):
```python
reshape_0 = gs.Node(name="reshape_0", op="Reshape")
reshape_0.inputs.append(input_new)
reshape_0_shape = gs.Constant(name="reshape_0_shape", values=np.array([MAX_VOXELS * 32, 10]))
reshape_0.inputs.append(reshape_0_shape)
# 输出:[V*32, 10]
```
- 插入线性层(第 107-120 行):
```python
# 找到图中原有的 MatMul 节点,修改其输入
matmul_op = [node for node in graph.nodes if node.op == "MatMul"][0]
matmul_op.inputs[0] = reshape_0_out # [V*32, 10]
# 插入 BatchNorm + ReLU
bn_op = [node for node in graph.nodes if node.op == "BatchNormalization"][0]
relu_op = [node for node in graph.nodes if node.op == "Relu"][0]
```
- 再次 Reshape(第 122-128 行):
```python
reshape_1 = gs.Node(name="reshape_1", op="Reshape")
reshape_1.inputs.append(relu_op_out)
reshape_1_shape = gs.Constant(name="reshape_1_shape", values=np.array([MAX_VOXELS, 32, 64]))
# 输出:[V, 32, 64]
```
- ReduceMax(第 130-132 行):
```python
reducemax_op = [node for node in graph.nodes if node.op == "ReduceMax"][0]
reducemax_op.inputs[0] = reshape_1_out
reducemax_op.attrs['keepdims'] = [0]
# 输出:[V, 64](每个体素一个 64 维特征)
```
- 替换为 PPScatterPlugin(第 134-135 行):
```python
conv_op = [node for node in graph.nodes if node.op == "Conv"][0]
graph.replace_with_clip(
[reducemax_op.outputs[0], X, Y], # 输入:pillar_features[V,64], voxel_idxs[V,4], voxel_num[1]
[conv_op.inputs[0]] # 输出:spatial_features
)
```
这里的 replace_with_clip 是在 replace_with_clip 函数中注册的自定义节点替换器(第 21-32 行):
```python
@gs.Graph.register()
def replace_with_clip(self, inputs, outputs):
op_attrs = dict()
op_attrs["dense_shape"] = np.array([496, 432]) # 输出 BEV 图像尺寸
return self.layer(
name="PPScatter_0",
op="PPScatterPlugin", # ← TensorRT 自定义插件算子
inputs=inputs,
outputs=outputs,
attrs=op_attrs
)
```
四、--plugins=build/libpointpillar_core.so 的作用
构建命令
trtexec --onnx=./model/pointpillar.onnx \
--fp16 \
--plugins=build/libpointpillar_core.so \
--saveEngine=./model/pointpillar.plan \
--inputIOFormats=fp16:chw,int32:chw,int32:chw \
--verbose \
--dumpLayerInfo \
--dumpProfile \
--separateProfileRun \
--profilingVerbosity=detailed
这个 .so 文件是什么?
从 CMakeLists.txt 可以看到编译内容:
file(GLOB_RECURSE CORE_FILES
src/pointpillar/*.cu # 所有 CUDA kernel
src/pointpillar/*.cpp # 所有 C++ 实现
src/common/tensor.cu
src/common/tensorrt.cpp
)
cuda_add_library(pointpillar_core SHARED ${CORE_FILES})
包含的模块:
| 源文件 | 功能 |
|----------|------|
| pointpillar-scatter.cpp | PPScatterPlugin 的 TensorRT 插件注册和实现 |
| pillarscatter-kernel.cu | Scatter 操作的 CUDA kernel(FP32/FP16) |
| lidar-backbone.cu/.cpp | 2D backbone 推理(Conv/Pooling/UpSampling) |
| lidar-postprocess.cu/.cpp | NMS 后处理 |
| lidar-voxelization.cu/.cpp | 点云体素化预处理 |
为什么需要 --plugins 参数?
TensorRT 原生只支持标准算子(如 Conv、MatMul 等)。当解析 ONNX 遇到 op="PPScatterPlugin" 这样的自定义算子时,TensorRT 不知道如何处理。
--plugins 参数的作用:
-
加载
libpointpillar_core.so共享库 -
调用库中的
REGISTER_TENSORRT_PLUGIN(PPScatterPluginCreator)(见pointpillar-scatter.cpp第 313 行) -
将
PPScatterPluginCreator注册到 TensorRT 的插件工厂 -
当解析器遇到
op="PPScatterPlugin"的节点时,调用PPScatterPluginCreator::createPlugin()
注册流程:
// pointpillar-scatter.cpp
// 1. 定义插件创建器
class PPScatterPluginCreator : public nvinfer1::IPluginCreator {
public:
// 根据属性创建插件实例
IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override {
const PluginField* fields = fc->fields;
int nbFields = fc->nbFields;
int target_h = 0, target_w = 0;
// 解析属性
for (int i = 0; i < nbFields; ++i) {
const char* attr_name = fields[i].name;
if (!strcmp(attr_name, "dense_shape")) {
const int* ts = static_cast<const int*>(fields[i].data);
target_h = ts[0]; // 496
target_w = ts[1]; // 432
}
}
// 创建插件实例
auto* plugin = new PPScatterPlugin(target_h, target_w);
return plugin;
}
};
// 2. 注册插件到 TensorRT(库加载时自动执行)
REGISTER_TENSORRT_PLUGIN(PPScatterPluginCreator);
PPScatterPlugin 具体做了什么?
从 pillarscatter-kernel.cu 可以看到,它执行的是 Pillar Scatter 操作。
算法原理:
输入:
pillar_features: [V, 64] ← 每个有效体素的 64 维特征
voxel_coords: [V, 4] ← 体素坐标 (batch, z, y, x)
voxel_num: [1] ← 有效体素数量
输出:
spatial_features: [1, 64, 496, 432] ← 稠密 BEV 伪图像
操作:
for i = 0 to voxel_num[0]-1:
y = voxel_coords[i, 2]
x = voxel_coords[i, 3]
spatial_features[0, :, y, x] = pillar_features[i, :]
CUDA Kernel 实现(FP16 版本):
__global__ void pillarScatterHalfKernel(
const half *pillar_features_data, // [V, 64]
const unsigned int *coords_data, // [V, 4]
const unsigned int *params_data, // [1] = voxel_num
unsigned int featureX, // 432
unsigned int featureY, // 496
half *spatial_feature_data // [1, 64, 496, 432]
) {
int pillar_idx = blockIdx.x * PILLARS_PER_BLOCK + threadIdx.x;
int num_pillars = params_data[0];
if (pillar_idx >= num_pillars) return;
// 读取体素坐标
uint4 coord = ((const uint4 *)coords_data)[pillar_idx];
unsigned int x = coord.w;
unsigned int y = coord.z;
// 读取体素特征到共享内存
__shared__ half pillarSM[PILLARS_PER_BLOCK][PILLAR_FEATURE_SIZE];
for (int i = 0; i < PILLAR_FEATURE_SIZE; i++) {
pillarSM[threadIdx.x][i] = pillar_features_data[pillar_idx * 64 + i];
}
__syncthreads();
// 写入 BEV 图像(HWC8 布局优化)
int C_stride = (64 + 7) / 8 * 8;
for (int i = 0; i < 64; i++) {
spatial_feature_data[y * featureX * C_stride + x * C_stride + i] =
pillarSM[threadIdx.x][i];
}
}
关键优化:
-
共享内存:每个 block 的 64 个 pillar 先加载到共享内存,减少全局内存访问
-
HWC8 布局(FP16):通道维度填充到 8 的倍数,利用 TensorRT 的向量化指令
-
动态批大小:根据
voxel_num[0]只处理有效体素,跳过空的占位体素
FP16 精度的支持
PPScatterPlugin 同时支持 FP32 和 FP16:
int PPScatterPlugin::enqueue(
const PluginTensorDesc* inputDesc,
const PluginTensorDesc* outputDesc,
const void* const* inputs,
void* const* outputs,
void* workspace,
cudaStream_t stream
) noexcept {
DataType inputType = inputDesc[0].type;
if (inputType == DataType::kHALF) {
// FP16 分支
auto pillar_features_data = static_cast<const half*>(inputs[0]);
auto spatial_feature_data = static_cast<half*>(outputs[0]);
cudaMemsetAsync(spatial_feature_data, 0, ..., stream);
return pillarScatterHalfKernelLaunch(...);
}
else if (inputType == DataType::kFLOAT) {
// FP32 分支
auto pillar_features_data = static_cast<const float*>(inputs[0]);
auto spatial_feature_data = static_cast<float*>(outputs[0]);
cudaMemsetAsync(spatial_feature_data, 0, ..., stream);
return pillarScatterFloatKernelLaunch(...);
}
}
配合 --fp16 参数,整个 TensorRT 引擎可以在 FP16 模式下运行:
-
计算吞吐量翻倍(FP16 是 FP32 的 2 倍速度)
-
显存占用减半
-
对精度影响可忽略(检测任务容忍度较高)
五、完整流程总结
训练到部署的完整 pipeline
┌────────────────────────────────────────────────────────────────────────┐
│ OpenPCDet 训练阶段 │
│ (Python + PyTorch) │
│ │
│ pointpillar.yaml → 模型定义 → 训练 → pointpillar_*.pth (权重) │
└────────────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────────────────────────────────┐
│ ONNX 导出阶段 (Python) │
│ │
│ 1. build_network() + load_params_from_file() → 恢复完整模型 │
│ 2. torch.onnx.export(dummy_input) → pointpillar_raw.onnx │
│ └─ 包含完整计算图(预处理+骨干+检测头+后处理) │
│ 3. simplify_postprocess() → 截断后处理 │
│ 4. onnxsim.simplify() → 常量折叠等优化 │
│ 5. simplify_preprocess() → 重构 PFE + 替换算子 │
│ └─ voxels: [V,32,4] → [V,32,10] │
│ └─ ScatterBEV → PPScatterPlugin │
│ │
│ 输出:pointpillar.onnx │
│ └─ 输入:voxels[V,32,10] + voxel_idxs[V,4] + voxel_num[1] │
│ └─ 输出:cls_preds[1,248,216,18] + box_preds + dir_cls_preds │
└────────────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────────────────────────────────┐
│ TensorRT 引擎构建 (C++) │
│ │
│ 1. 编译 libpointpillar_core.so │
│ ├─ PPScatterPlugin(TensorRT 插件注册) │
│ ├─ pillarScatterKernel(CUDA kernel) │
│ ├─ 2D Backbone(Conv/Pooling) │
│ └─ PostProcess(NMS) │
│ │
│ 2. trtexec --onnx=pointpillar.onnx │
│ --fp16 │
│ --plugins=libpointpillar_core.so ← 加载自定义算子 │
│ --saveEngine=pointpillar.plan │
│ │
│ 输出:pointpillar.plan(TensorRT 序列化引擎) │
└────────────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────────────────────────────────┐
│ 推理阶段 (C++) │
│ │
│ 原始点云 (.bin/.pcd) │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────────────┐ │
│ │ CUDA Voxelization (lidar-voxelization.cu) │ │
│ │ ├─ 点云 → 体素网格 │ │
│ │ ├─ 计算额外特征(偏移量等) │ │
│ │ └─ 输出 voxels[V,32,10] + voxel_idxs[V,4] + voxel_num[1] │ │
│ └────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────────────┐ │
│ │ TensorRT Engine (pointpillar.plan) │ │
│ │ ├─ PillarEncoder → PPScatterPlugin │ │
│ │ │ └─ [V,64] → [1,64,496,432] (BEV 伪图像) │ │
│ │ ├─ 2D Backbone (Conv/Pooling/UpSampling) │ │
│ │ └─ Detect Head (分类/框回归/方向) │ │
│ └────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ cls_preds[1,248,216,18] box_preds[1,248,216,42] dir_cls_preds │
│ │ │
│ ▼ │
│ ┌────────────────────────────────────────────────────────────┐ │
│ │ NMS PostProcess (lidar-postprocess.cu) │ │
│ │ ├─ 解码框回归参数 │ │
│ │ ├─ 过滤低置信度框 │ │
│ │ └─ 非极大值抑制 │ │
│ └────────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ 最终检测框(3D 边界框 + 类别 + 置信度) │
└────────────────────────────────────────────────────────────────────────┘
关键文件对照表
| 阶段 | 文件 | 功能 |
|------|------|------|
| 训练 | pointpillar.yaml | 模型配置(anchor 尺寸、检测类别等) |
| 训练 | *.pth | OpenPCDet 训练权重 |
| 导出 | tool/export_onnx.py | 主导出脚本 |
| 导出 | tool/modify_onnx.py | ONNX 图修改(预处理重构、后处理截断) |
| 构建 | CMakeLists.txt | 编译配置 |
| 构建 | src/pointpillar/pointpillar-scatter.cpp | PPScatterPlugin TensorRT 插件 |
| 构建 | src/pointpillar/pillarscatter-kernel.cu | Scatter CUDA kernel |
| 构建 | tool/build_trt_engine.sh | TensorRT 引擎构建脚本 |
| 推理 | src/pointpillar/lidar-voxelization.cu | 体素化预处理 |
| 推理 | src/pointpillar/lidar-backbone.cu | 2D backbone 推理 |
| 推理 | src/pointpillar/lidar-postprocess.cu | NMS 后处理 |
| 推理 | src/main.cpp | 主程序入口 |
性能优化要点
-
FP16 精度:
--fp16+ PPScatterPlugin 的 half kernel 实现,吞吐量翻倍 -
稀疏优化:只处理有效体素(
voxel_num),减少空洞计算 -
共享内存:PillarScatterKernel 使用 shared memory 减少全局内存访问
-
HWC8 布局:FP16 分支使用向量化内存布局
-
常量折叠:ONNX simplify 阶段合并常量算子,减少节点数
-
自定义算子:PPScatterPlugin 替换通用 Scatter,针对 BEV 特定布局优化
六、自定义训练模型适配指南
本节说明如何将上述标准ONNX导出流程适配到自定义训练的PointPillar模型。以container检测模型为例,展示关键参数的适配方法。
6.1 配置差异分析
6.1.1 关键配置参数对比
| 参数 | 标准KITTI模型 | 自定义container模型 | 影响 |
|------|---------------|-------------------|------|
| 检测类别 | 3类 | 1类 | 输出维度、锚框数量 |
| POINT_CLOUD_RANGE | [0, -40.0, -3, 70.4, 40.0, 1] | [-7.23, -33.07, -5.71, 75.20, 39.73, 1.02] | BEV特征图尺寸 |
| VOXEL_SIZE | [0.16, 0.16, 4] | [0.1, 0.1, 6.73] | 体素分辨率、内存占用 |
| MAX_NUMBER_OF_VOXELS | 10000-12000 | 150000 | TensorRT输入尺寸 |
| 锚框数量 | 每类6个anchor,共18个 | 每类3个anchor(仅1类),共3个 | 输出通道维度 |
6.1.2 自定义配置示例
# pointpillar.yaml(自定义container检测)
CLASS_NAMES: ['container']
DATA_CONFIG:
POINT_CLOUD_RANGE: [-7.23, -33.07, -5.71, 75.20, 39.73, 1.02]
DATA_PROCESSOR:
- NAME: transform_points_to_voxels
VOXEL_SIZE: [0.1, 0.1, 6.73]
MAX_POINTS_PER_VOXEL: 32
MAX_NUMBER_OF_VOXELS: 150000
MODEL:
VFE:
NUM_FILTERS: [64] # 柱子特征维度
MAP_TO_BEV:
NUM_BEV_FEATURES: 64
BACKBONE_2D:
LAYER_NUMS: [3, 5, 5]
LAYER_STRIDES: [2, 2, 2]
NUM_FILTERS: [64, 128, 256]
UPSAMPLE_STRIDES: [1, 2, 4]
NUM_UPSAMPLE_FILTERS: [128, 128, 128]
DENSE_HEAD:
ANCHOR_GENERATOR_CONFIG: [
{
'class_name': 'container',
'anchor_sizes': [[12.40, 2.57, 2.87]], # 仅1个anchor尺寸
'anchor_rotations': [0, 1.57, 3.05], # 3个方向
'anchor_bottom_heights': [-3.6],
'feature_map_stride': 2,
}
]
6.2 输出张量维度计算
6.2.1 BEV特征图尺寸
根据点云范围和体素大小计算:
BEV宽度 = floor((X_max - X_min) / voxel_size_x)
BEV高度 = floor((Y_max - Y_min) / voxel_size_y)
标准KITTI模型:
X范围: 0 → 70.4, Y范围: -40.0 → 40.0
VOXEL_SIZE: [0.16, 0.16]
BEV宽度 = 70.4 / 0.16 = 440
BEV高度 = 80.0 / 0.16 = 500
经过backbone下采样(stride=2)后:
特征图尺寸: [220, 250] ≈ [248, 216](考虑padding)
自定义container模型:
X范围: -7.23 → 75.20 (跨度82.43)
Y范围: -33.07 → 39.73 (跨度72.80)
VOXEL_SIZE: [0.1, 0.1]
BEV宽度 = 82.43 / 0.1 = 824
BEV高度 = 72.80 / 0.1 = 728
经过backbone下采样(stride=2)后:
特征图尺寸: [412, 364]
6.2.2 输出张量维度
标准KITTI模型(3类 × 6anchor = 18个锚框):
| 输出张量 | 形状 | 维度说明 |
|---------|------|---------|
| cls_preds | [1, 248, 216, 18] | 每个位置18个anchor的分类得分 |
| box_preds | [1, 248, 216, 42] | 每个anchor 7维编码(3个方向 × 7维 = 21,但实际42维) |
| dir_cls_preds | [1, 248, 216, 12] | 每个anchor 2个方向分类(18个 × 2 = 36?实际12维) |
自定义container模型(1类 × 3anchor = 3个锚框):
| 输出张量 | 形状 | 维度说明 |
|---------|------|---------|
| cls_preds | [1, 364, 412, 3] | 每个位置3个anchor的分类得分 |
| box_preds | [1, 364, 412, 21] | 每个anchor 7维编码(3个anchor × 7维 = 21) |
| dir_cls_preds | [1, 364, 412, 6] | 每个anchor 2个方向分类(3个anchor × 2 = 6) |
6.3 export_onnx.py 适配修改
6.3.1 修改 MAX_VOXELS 参数
第158行(硬编码的MAX_VOXELS需要根据配置修改):
# 原始代码
MAX_VOXELS = 150000 # ← 硬编码,需要从配置读取
# 建议修改为
MAX_VOXELS = cfg.DATA_CONFIG.DATA_PROCESSOR[2].MAX_NUMBER_OF_VOXELS['test']
或者直接读取配置:
# 找到 transform_points_to_voxels 的配置
for processor in cfg.DATA_CONFIG.DATA_PROCESSOR:
if processor.NAME == 'transform_points_to_voxels':
MAX_VOXELS = processor.MAX_NUMBER_OF_VOXELS['test']
VOXEL_SIZE = processor.VOXEL_SIZE
break
6.3.2 命令行参数适配
# 标准KITTI模型
python tool/export_onnx.py \
--cfg_file cfgs/kitti_models/pointpillar.yaml \
--ckpt pointpillar_7768.pth \
--data_path data/kitti/velodyne/000000.bin \
--out_dir model/
# 自定义container模型
python tool/export_onnx.py \
--cfg_file CUDA-PointPillars/pointpillar.yaml \
--ckpt custom_container_checkpoint.pth \
--data_path custom_data/test.bin \
--out_dir model/
6.4 modify_onnx.py 适配修改
6.4.1 修改输出张量形状(simplify_postprocess函数)
第44-46行需要根据实际模型输出修改:
# 原始代码(KITTI标准配置)
cls_preds = gs.Variable(name="cls_preds", dtype=np.float32, shape=(1, 248, 216, 18))
box_preds = gs.Variable(name="box_preds", dtype=np.float32, shape=(1, 248, 216, 42))
dir_cls_preds = gs.Variable(name="dir_cls_preds", dtype=np.float32, shape=(1, 248, 216, 12))
# 自定义container模型适配
cls_preds = gs.Variable(name="cls_preds", dtype=np.float32, shape=(1, 364, 412, 3))
box_preds = gs.Variable(name="box_preds", dtype=np.float32, shape=(1, 364, 412, 21))
dir_cls_preds = gs.Variable(name="dir_cls_preds", dtype=np.float32, shape=(1, 364, 412, 6))
如何自动计算输出形状?
可以通过解析配置文件自动计算:
def compute_output_shape(cfg):
"""
根据 pointpillar.yaml 配置计算ONNX输出形状
"""
# 1. 读取点云范围和体素大小
pc_range = cfg.DATA_CONFIG.POINT_CLOUD_RANGE
for processor in cfg.DATA_CONFIG.DATA_PROCESSOR:
if processor.NAME == 'transform_points_to_voxels':
voxel_size = processor.VOXEL_SIZE
break
# 2. 计算BEV尺寸
bev_width = int((pc_range[3] - pc_range[0]) / voxel_size[0])
bev_height = int((pc_range[4] - pc_range[1]) / voxel_size[1])
# 3. 考虑backbone下采样(stride=2)
feat_width = bev_width // 2
feat_height = bev_height // 2
# 4. 计算anchor数量
num_classes = len(cfg.CLASS_NAMES)
anchors_per_class = 0
for anchor_cfg in cfg.MODEL.DENSE_HEAD.ANCHOR_GENERATOR_CONFIG:
anchors_per_class += len(anchor_cfg['anchor_rotations'])
num_anchors = num_classes * anchors_per_class
# 5. 计算输出维度
num_dir_bins = cfg.MODEL.DENSE_HEAD.NUM_DIR_BINS
cls_dim = num_anchors
box_dim = num_anchors * 7 # (x, y, z, w, l, h, theta)
dir_dim = num_anchors * num_dir_bins
return {
'cls_shape': (1, feat_height, feat_width, cls_dim),
'box_shape': (1, feat_height, feat_width, box_dim),
'dir_shape': (1, feat_height, feat_width, dir_dim)
}
使用方法:
# 在 simplify_postprocess 函数开头
output_shapes = compute_output_shape(cfg)
cls_preds = gs.Variable(name="cls_preds", dtype=np.float32, shape=output_shapes['cls_shape'])
box_preds = gs.Variable(name="box_preds", dtype=np.float32, shape=output_shapes['box_shape'])
dir_cls_preds = gs.Variable(name="dir_cls_preds", dtype=np.float32, shape=output_shapes['dir_shape'])
6.4.2 修改 PPScatterPlugin 的 dense_shape(replace_with_clip函数)
第30行的 dense_shape 需要根据BEV尺寸修改:
# 原始代码(KITTI标准配置)
op_attrs["dense_shape"] = np.array([496, 432]) # BEV输出尺寸
# 自定义container模型适配
op_attrs["dense_shape"] = np.array([728, 824]) # BEV输出尺寸
如何自动计算 dense_shape?
# 在 replace_with_clip 函数中根据配置计算
def compute_dense_shape(cfg):
pc_range = cfg.DATA_CONFIG.POINT_CLOUD_RANGE
for processor in cfg.DATA_CONFIG.DATA_PROCESSOR:
if processor.NAME == 'transform_points_to_voxels':
voxel_size = processor.VOXEL_SIZE
break
dense_height = int((pc_range[4] - pc_range[1]) / voxel_size[1])
dense_width = int((pc_range[3] - pc_range[0]) / voxel_size[0])
return np.array([dense_height, dense_width])
# 在注册函数中使用
@gs.Graph.register()
def replace_with_clip(self, inputs, outputs):
for inp in inputs:
inp.outputs.clear()
for out in outputs:
out.inputs.clear()
op_attrs = dict()
op_attrs["dense_shape"] = compute_dense_shape(cfg) # ← 动态计算
return self.layer(name="PPScatter_0", op="PPScatterPlugin", inputs=inputs, outputs=outputs, attrs=op_attrs)
6.5 TensorRT 引擎构建适配
6.5.1 输入/输出尺寸更新
TensorRT 构建(trtexec)需要指定正确的输入输出格式:
# 标准KITTI模型
trtexec --onnx=./model/pointpillar.onnx \
--fp16 \
--plugins=build/libpointpillar_core.so \
--saveEngine=./model/pointpillar.plan \
--inputIOFormats=fp16:chw,int32:chw,int32:chw \
--minShapes=voxels:1x32x10,voxel_num:1x1,voxel_idxs:1x4 \
--optShapes=voxels:10000x32x10,voxel_num:1x1,voxel_idxs:10000x4 \
--maxShapes=voxels:12000x32x10,voxel_num:1x1,voxel_idxs:12000x4
# 自定义container模型
trtexec --onnx=./model/pointpillar.onnx \
--fp16 \
--plugins=build/libpointpillar_core.so \
--saveEngine=./model/pointpillar.plan \
--inputIOFormats=fp16:chw,int32:chw,int32:chw \
--minShapes=voxels:1x32x10,voxel_num:1x1,voxel_idxs:1x4 \
--optShapes=voxels:150000x32x10,voxel_num:1x1,voxel_idxs:150000x4 \
--maxShapes=voxels:150000x32x10,voxel_num:1x1,voxel_idxs:150000x4
6.5.2 后处理参数适配
TensorRT 引擎输出后,后处理需要适配新的输出维度:
// C++ 推理代码适配示例
// 原始代码(KITTI标准配置)
const int cls_channels = 18; // 3类 × 6anchor
const int box_channels = 42; // 18anchor × 7维
const int dir_channels = 12; // 18anchor × 2方向
// 自定义container模型
const int cls_channels = 3; // 1类 × 3anchor
const int box_channels = 21; // 3anchor × 7维
const int dir_channels = 6; // 3anchor × 2方向
// 后处理NMS时也需要相应调整
const int num_classes = 1; // container
const int num_anchors = 3; // 每类3个anchor
const float score_thresh = 0.1; // 从配置文件读取
const float nms_thresh = 0.1; // 从配置文件读取
6.6 完整适配流程总结
Step 1: 修改配置文件
确保 pointpillar.yaml 包含正确的自定义配置:
CLASS_NAMES: ['container']
POINT_CLOUD_RANGE: [x_min, y_min, z_min, x_max, y_max, z_max]
VOXEL_SIZE: [vx, vy, vz]
MAX_NUMBER_OF_VOXELS: 150000
ANCHOR_GENERATOR_CONFIG: [自定义锚框配置...]
Step 2: 修改 export_onnx.py
# 添加配置读取逻辑
MAX_VOXELS = 从cfg.DATA_CONFIG.DATA_PROCESSOR读取
Step 3: 修改 modify_onnx.py
# 在 simplify_postprocess 中修改输出形状
cls_preds = (1, H, W, num_anchors) # H/W 根据BEV尺寸计算
box_preds = (1, H, W, num_anchors * 7)
dir_cls_preds = (1, H, W, num_anchors * 2)
# 在 replace_with_clip 中修改 dense_shape
op_attrs["dense_shape"] = (bev_height, bev_width)
Step 4: 导出ONNX
python tool/export_onnx.py \
--cfg_file your_custom_config.yaml \
--ckpt your_custom_checkpoint.pth \
--data_path sample_data.bin \
--out_dir model/
Step 5: 构建TensorRT引擎
trtexec --onnx=./model/pointpillar.onnx \
--fp16 \
--plugins=build/libpointpillar_core.so \
--saveEngine=./model/pointpillar.plan \
--optShapes=voxels:150000x32x10,voxel_num:1x1,voxel_idxs:150000x4
Step 6: 适配C++推理代码
// 修改输出张量维度和后处理参数
const int cls_channels = 3;
const int box_channels = 21;
const int dir_channels = 6;
const float score_thresh = 从配置读取;
const float nms_thresh = 从配置读取;
6.7 常见问题排查
问题1: ONNX 导出失败,提示维度不匹配
原因: MAX_VOXELS 与实际数据不一致
解决: 检查 pointpillar.yaml 中的 MAX_NUMBER_OF_VOXELS,确保与 export_onnx.py 中的 MAX_VOXELS 一致
问题2: TensorRT 构建失败,提示输入/输出尺寸错误
原因: modify_onnx.py 中的输出形状与实际模型不符
解决: 使用 compute_output_shape() 函数自动计算,或手动检查配置文件中的锚框数量和点云范围
问题3: 推理结果精度下降
原因: BEV分辨率或锚框尺寸与训练时配置不一致
解决: 确保推理时的 VOXEL_SIZE、POINT_CLOUD_RANGE 和 ANCHOR_GENERATOR_CONFIG 与训练配置完全一致
问题4: 内存不足(OOM)
原因: MAX_VOXELS 设置过大(如从10000改到150000)
解决:
-
调整
MAX_NUMBER_OF_VOXELS到合理范围(如50000-80000) -
使用
--minShapes和--maxShapes设置动态范围 -
确保
--fp16参数启用
6.8 适配检查清单
在将自定义模型部署到TensorRT前,逐项检查:
附录:技术术语对照表
| 术语 | 英文 | 解释 |
|------|------|------|
| 体素 | Voxel | 3D 空间中的离散立方体单元,用于表示点云密度 |
| 柱子 | Pillar | Z 轴方向无限延伸的体素,PointPillar 的核心概念 |
| 散射 | Scatter | 将稀疏数据映射到稠密空间的操作(类似稀疏矩阵转稠密) |
| BEV | Bird's Eye View | 俯视图,将 3D 点云投影到 2D 平面的视角 |
| Backbone | 骨干网络 | 提取特征的主干卷积网络 |
| PFE | Pillar Feature Encoder | 柱子特征编码器 |
| VFE | Voxel Feature Encoder | 体素特征编码器 |
| NMS | Non-Maximum Suppression | 非极大值抑制,去除重复检测框 |
| Opset | ONNX 算子集版本 | ONNX 标准的版本号,规定了支持的算子集合 |
| Plugin | 插件 | TensorRT 的扩展机制,允许用户注册自定义算子 |
| TensorRT Plan | TensorRT 序列化引擎 | 编译优化后的推理引擎文件 |
适用版本:CUDA-PointPillars (NVIDIA官方实现) + openpcdet 0.6.0 (latest)

浙公网安备 33010602011771号