cuda-pointpillar+openpcdet 0.6.0 自定义模型适配指北

目录


一、整体架构概览

整个推理 pipeline 分为 4 个阶段,但 ONNX 只导出中间 2 个:


┌──────────────┐    ┌──────────────────┐    ┌──────────────────┐    ┌──────────────┐

│ Voxelization │ →  │  PillarEncoder   │ →  │   2D Backbone    │ →  │  PostProcess │

│  (预处理,CPU) │    │  + ScatterBEV    │    │  + DetectHead    │    │  (NMS,CPU)   │

│              │    │  ← ONNX 导出这段 →│    │  ← ONNX 导出这段 →│    │              │

└──────────────┘    └──────────────────┘    └──────────────────┘    └──────────────┘

阶段说明

  1. Voxelization:将原始 3D 点云转换为体素(voxel)表示,由 CUDA/C++ 预处理模块完成,不在 ONNX 中

  2. PillarEncoder + ScatterBEV:提取每个柱子(pillar)的特征,散射到 2D 伪图像,ONNX 导出

  3. 2D Backbone + DetectHead:2D 卷积骨干网络 + 检测头输出分类/框回归/方向预测,ONNX 导出

  4. 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 到检测头输出的部分。

详细步骤

  1. 定义最终输出节点(第 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))

```

  1. 定位关键节点(第 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]]

```

  1. 清理输入/输出(第 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

关键代码解析

  1. 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]

```

  1. 插入线性层(第 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]

```

  1. 再次 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]

```

  1. 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 维特征)

```

  1. 替换为 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 参数的作用:

  1. 加载 libpointpillar_core.so 共享库

  2. 调用库中的 REGISTER_TENSORRT_PLUGIN(PPScatterPluginCreator)(见 pointpillar-scatter.cpp 第 313 行)

  3. PPScatterPluginCreator 注册到 TensorRT 的插件工厂

  4. 当解析器遇到 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 | 主程序入口 |

性能优化要点

  1. FP16 精度--fp16 + PPScatterPlugin 的 half kernel 实现,吞吐量翻倍

  2. 稀疏优化:只处理有效体素(voxel_num),减少空洞计算

  3. 共享内存:PillarScatterKernel 使用 shared memory 减少全局内存访问

  4. HWC8 布局:FP16 分支使用向量化内存布局

  5. 常量折叠:ONNX simplify 阶段合并常量算子,减少节点数

  6. 自定义算子: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_SIZEPOINT_CLOUD_RANGEANCHOR_GENERATOR_CONFIG 与训练配置完全一致

问题4: 内存不足(OOM)

原因: MAX_VOXELS 设置过大(如从10000改到150000)

解决:

  1. 调整 MAX_NUMBER_OF_VOXELS 到合理范围(如50000-80000)

  2. 使用 --minShapes--maxShapes 设置动态范围

  3. 确保 --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)

posted @ 2026-06-01 15:31  Yihoyo  阅读(35)  评论(0)    收藏  举报