PointPillars 点云目标检测实战：TensorFlow Model Garden 中的 BEV 编码、训练流水线与配置详解

本文以 TensorFlow Model Garden（official/projects/pointpillars）中的 PointPillars 实现为主线，系统讲解该模型如何把原始 3D 点云编码为鸟瞰图（Bird-Eye-View, BEV）伪图像、再交给标准 2D 卷积检测架构的完整原理与工程实现。读完本文，你将掌握：PointPillars 各网络模块（Featurizer / Backbone / Decoder / SSDHead / DetectionGenerator）的源码级结构与张量形状约定、Waymo Open Dataset 的 Beam 预处理流程、train.py 训练入口的调用链，以及 GPU/TPU 两套基线 YAML 配置中每一项参数的含义与调优思路。

1. PointPillars 是什么：从原始点云到 BEV 图像

PointPillars 是一个点云目标检测模型（出自论文 arXiv:1812.05784）。它的核心思想是：把原始 3D 点云信号编码成一种适合下游检测流水线的格式——鸟瞰图（BEV image）。具体做法是：

1. 沿竖直方向把点云划分为若干"立柱"（pillars），每一根立柱收集落在同一水平网格单元上的所有点；
2. 用类似 PointNet 的编码器学习每根立柱内点集的表示；
3. 将编码后的立柱特征"散点"（scatter）回一张 BEV 伪图像上，使其可以输入任意标准 2D 卷积检测网络。

本仓库的实现基于 TensorFlow Model Garden 的通用训练框架，官方 README（official/projects/pointpillars/README.md）给出的关键指标是：在 Waymo Open Dataset 1.2.0 上训练，vehicle 类别达到 45.96% mAP / 45.35% mAPH，单张 V100 GPU、batch size 1 的推理时间为 53ms。

整个项目目录组织如下，每个子目录对应流水线的一个环节：

目录	职责
configs/	实验配置定义（pointpillars_baseline 实验类型）与三份基线 YAML
dataloaders/	TFRecord 解码器与标签标注器（parser）
modeling/	模型五大组件：Featurizer、Backbone、Decoder、Head、整体 Model 与工厂函数
tasks/	训练/验证步骤、损失函数与 Waymo 评测指标
tools/	Waymo 数据预处理脚本与模型导出工具
utils/	锚点生成、检测评估器、SavedModel 导出等辅助模块

2. 环境与依赖安装

README 给出的环境要求是 TensorFlow 2.6 生态（这也是 tf_keras 别名在仓库代码中广泛使用的原因）：

pip install --upgrade pip
pip install tensorflow==2.6.0
pip install tf-models-official==2.7.2
pip install apache-beam[gcp]==2.42.0 --user

其中 apache-beam 只用于数据集预处理脚本（tools/process_wod.py 基于 Beam pipeline 读写 TFRecord）；若只在已有处理后数据的机器上训练，可以不装 GCP 扩展。

3. 数据集准备：Waymo 原始数据 → 模型可吃的 TFRecord

3.1 安装 Waymo 官方库

以 Waymo Open Dataset 为例，需要先安装与 TensorFlow 版本严格匹配的预编译包（README 指定的是 2.6 版本）：

pip install waymo-open-dataset-tf-2-6-0

注意：该 pip 包是针对特定 TF 版本构建的，与当前安装的 TF 版本不一致会报错。这一点在任务代码中有明确注释（tasks/pointpillars.py 的 build_metrics）。

3.2 运行转换脚本

使用仓库提供的 tools/process_wod.py，把 Waymo 原始 lidar frame（dataset.proto 中的 Frame）转成模型可直接读取的 tf.Example 序列：

SRC_DIR="gs://waymo_open_dataset_v_1_2_0_individual_files"
DST_DIR="gs://<path/to/directory>"
# 分布式 runner 参考 Apache Beam 官方文档
RUNNER="DirectRunner"

python3 process_wod.py \
--src_dir=${SRC_DIR} \
--dst_dir=${DST_DIR} \
--pipeline_options="--runner=${RUNNER}"

从源码可以看到该脚本的实现细节（tools/process_wod.py）：

• 入参：--src_dir（原始 WOD TFRecord 目录）、--dst_dir（输出目录）、--config_file（YAML 配置）、--pipeline_options（Beam runner 选项）；
• 硬性约束：--src_dir 下必须存在 training 与 validation 两个子目录（源码中 _SRC_FOLDERS = ['training', 'validation']，L46-L47），脚本会分别读取这两路数据；
• 输出格式：通过 tfrecordio.WriteToTFRecord 写出 gzip 压缩的 .tfrecord 文件（L69-L75），与训练端 DataConfig.file_type='tfrecord_compressed' 的默认值严格对应；
• 同时会写一份 .stats.txt 全局计数文件（count_examples），方便核对样本数。

转换的几何逻辑（立柱划分、索引计算）由 utils/wod_processor.py 中的 WodProcessor 完成，划分参数与训练配置中的 task.model.image / task.model.pillars 保持一致。

3.3 处理后每个样本包含哪些字段

解码端定义了完整的特征 schema（dataloaders/decoders.py）：

self._feature_description = {
    'frame_id': tf.io.FixedLenFeature([], tf.int64),
    'pillars': tf.io.FixedLenFeature([], tf.string),
    'indices': tf.io.FixedLenFeature([], tf.string),
    'bbox/ymin': tf.io.VarLenFeature(tf.float32),
    'bbox/xmin': tf.io.VarLenFeature(tf.float32),
    'bbox/ymax': tf.io.VarLenFeature(tf.float32),
    'bbox/xmax': tf.io.VarLenFeature(tf.float32),
    'bbox/class': tf.io.VarLenFeature(tf.int64),
    'bbox/heading': tf.io.VarLenFeature(tf.float32),
    'bbox/z': tf.io.VarLenFeature(tf.float32),
    'bbox/height': tf.io.VarLenFeature(tf.float32),
    'bbox/difficulty': tf.io.VarLenFeature(tf.int64),
}

其中 pillars 是原始字节串，解码时 reshape 为 [P, N, D]（P=立柱数、N=每立柱点数、D=每点特征维数）；indices 解码为 [P, 2] 的 int32 坐标（立柱在 BEV 网格中的行列位置）；bbox/* 系列是变长标签，解码后堆叠为 [M, 4] 的 yxyx 框（decoders.py 的 _decode_pillars / _decode_boxes）。

4. 模型结构：源码级拆解

模型由五个可序列化（register_keras_serializable）的组件构成，工厂函数 modeling/factory.py 的 build_pointpillars 按如下顺序组装：

pillars [B,P,N,D] + indices [B,P,2]
        │
        ▼  Featurizer（1x1 卷积块 + 柱内 max-pool + scatter_nd）
   BEV 图像 [B, 512, 512, 64]
        │
        ▼  Backbone（多级下采样 ConvBlock 组）
   {level: [B, 512/2^l, 512/2^l, 64·2^(l-1)]}
        │
        ▼  Decoder（上采样，把最深层上采回到各输出层分辨率）
   {level: 特征图}
        │
        ▼  SSDHead（共享权重的分类/回归/属性卷积）
   scores / boxes / attributes
        │
        ▼  DetectionGenerator（NMS 生成最终检测框，仅评估/导出时）
   boxes, scores, classes, attributes

4.1 Featurizer：PointNet 式立柱编码 + 散点回 BEV

实现见 modeling/featurizers.py。前向过程（call）为：

1. 给每根立柱的网格坐标 indices（[B, P, 2]）拼上 batch 维，得到 batch_indices（[B, P, 3]）；
2. 依次通过 num_blocks 个 ConvBlock（kernel_size=1 的 1x1 卷积，等价于对点做逐点 MLP），[B, P, N, D] → [B, P, N, C]；
3. 沿点维做 max pooling（tf.reduce_max(x, axis=2)）得到每根立柱的固定表示 [B, P, C]——这是 PointNet 思想的直接体现：置换不变性来自对称函数 max；
4. 用 tf.scatter_nd 把立柱特征写回 BEV 网格，得到 [B, H, W, C] 的伪图像。

一个工程细节值得注意：tf.scatter_nd 要求具体（concrete）batch size，所以 build 阶段为训练/评估/推理三种模式分别预构建了 batch_dims 张量（_get_batch_dims，L97-L110），并在 _get_batch_size_and_dims 中用 training 参数做三态区分（True/False/None），None 表示 SavedModel 导出时的 test 模式，此时固定 batch=1。

4.2 Backbone：多级下采样

实现见 modeling/backbones.py。它由 max_level 个"下采样组"串联而成，每组结构为：

• 首个 ConvBlock：3x3 卷积、strides=2（分辨率减半）；
• 后续 num_convs - 1 个 ConvBlock：3x3 卷积、strides=1；
• 通道数逐层翻倍（filters = input_channels * scale，scale 每级 ×2）。

第 l 级输出分辨率为 输入分辨率 / 2^l（L44-L47 注释），即 512×512 → 256×256 → 128×128（level 1/2/3）。min_level 限制为 ≥ 1（L70-L73），只从下采样后的特征开始输出。基线配置中 min_level=1, max_level=3, num_convs=6，对应论文中 VGG 风格的 6 层卷积堆叠。

4.3 Decoder 与 SSDHead

Decoder 接收 Backbone 的多级特征并上采样对齐（modeling/decoders.py），使所有输出层处于同一分辨率，这与 SSD 多尺度解码的设计一致。

检测头 modeling/heads.py 是一个 SSDHead：

• 分类分支 self._classifier：3x3 Conv2D，输出 num_classes * num_anchors_per_location 个通道，bias 初始化为 -log((1-0.01)/0.01) ≈ -4.6（L102），即初始预测正类概率约 1%，这是检测头常见的偏置初始化技巧；
• 框回归分支 self._box_regressor：输出 num_params_per_anchor(=4) * num_anchors_per_location 通道；
• 属性分支：为每个 attribute head（默认 heading、height、z 三个回归头，见 configs/pointpillars.py）各建一个 3x3 Conv2D；
• 所有分支在多级特征间共享权重。

最终输出由 modeling/models.py 的 PointPillarsModel.call 汇总：训练时（training=True）只返回 cls_outputs / box_outputs / attribute_outputs 三个多尺度字典；非训练时额外传入 image_shape 与 anchor_boxes，经 DetectionGenerator 做 NMS 后输出 boxes [B, M, 4]、scores [B, M]、classes [B, M]、attributes 与 num_detections。注意 generate_outputs 会把 heading 截断到 [-pi, pi]（L73-L75），并把所有原始预测 cast 到 float32 以支持混合精度训练。

5. 训练入口与调用链

训练入口 train.py 的 main 流程是 Model Garden 标准的六步：

1. gin.parse_config_files_and_bindings 解析 gin 绑定；
2. train_utils.parse_configuration(FLAGS) 解析实验类型 + YAML + params_override，得到完整 ExperimentConfig；
3. _check_if_resumed_job：检测 model_dir 中是否已有 checkpoint，判断任务是"断点续训"还是"全新训练"（L38-L55）。注释说明了 Cloud TPU 作业可能被机器调度器随时终止/恢复，续训作业会自动从 model_dir 恢复并跳过已完成的 step；
4. 按 runtime.mixed_precision_dtype 设置混合精度策略，并按 runtime.distribution_strategy（mirrored / tpu）构建 DistributionStrategy；
5. task_factory.get_task 按注册的 PointPillarsTask 构建任务对象，train_lib.run_experiment 跑训练/评估循环；
6. model_exporter.export_inference_graph（L96-L100）：训练结束后自动把 batch=1 的推理图导出到 <model_dir>/saved_model。

任务类 tasks/pointpillars.py 的关键方法：

• build_model：构造输入 spec（pillars: (None, P, N, D)、indices: (None, P, 2)，L76-L85），并用 get_batch_size_per_replica 把全局 batch size 均摊到各副本——要求 global_batch_size 必须能被副本数整除，否则直接抛错（L43-L55）。这就是为什么 GPU 配置写 16 # 2 * 8、TPU 配置写 64 # 2 * 32：全局 batch = 每副本 batch × 副本数；
• build_inputs：实例化 ExampleDecoder（只解码）与 Parser（负责把 GT 框按 match_threshold / unmatched_threshold 匹配到锚点上，生成 cls_targets / box_targets / attribute_targets 及正样本权重），再交给 input_reader_factory.input_reader_generator 组装 tf.data.Dataset（L131-L166）；
• train_step：前向 → compute_losses → 除以 num_replicas 得到 per-replica 损失 → 反向，兼容 LossScaleOptimizer 的缩放/还原（L299-L336）；
• validation_step / aggregate_logs / reduce_aggregated_logs：评估时逐 step 把 (groundtruths, outputs) 喂给 Waymo 检测评估器（wod_detection_evaluator.create_evaluator），在评估周期结束时聚合出 mAP/mAPH。

5.1 损失函数设计

损失实现（tasks/pointpillars.py）与配置类 Losses（configs/pointpillars.py）的对应关系：

损失项	类型	配置默认值	说明
class_loss	FocalLoss	focal_loss_alpha=0.25, focal_loss_gamma=1.5	缓解检测中极端的正负样本失衡
box_loss	Huber	huber_loss_delta=0.1	框回归，权重 box_loss_weight=100
attribute_loss	Huber	同上 delta	权重 attribute_loss_weight=10；heading 用方向感知处理：先把角度差 wrap 到 [-pi, pi]（utils.wrap_angle_rad）再算损失，避免 0 与 2pi 被判为巨大误差（L188-L196）

归一化方式也值得留意：分类/回归样本权重都除以 num_positives（批内正样本总数 + 1 防止 inf），因此每步损失量级与批内正样本数解耦（L223-L229）。

5.2 实验配置定义

pointpillars_baseline 实验类型在 configs/pointpillars.py 中通过 @exp_factory.register_config_factory('pointpillars_baseline') 注册，这正是训练命令 --experiment="pointpillars_baseline" 能解析到的原因。几个关键配置 dataclass：

• ImageConfig（L26-L41）：x_range / y_range / z_range / resolution 定义 BEV 覆盖范围与分辨率，height 和 width 不是手填的——__post_init__ 会自动按 (range / resolution) 计算，即 (-(-76.8) + 76.8) / 0.3 = 512；
• PillarsConfig（L44-L49）：num_pillars=24000、num_points_per_pillar=100、num_features_per_point=10，与预处理端必须严格一致，否则解码 reshape 会失败；
• AnchorLabeler（L82-L86）：match_threshold / unmatched_threshold 控制 GT 与锚点做正负样本匹配的 IoU 阈值（baseline 用 0.6 / 0.45）；
• DetectionGenerator（L129-L138）：pre_nms_top_k、pre_nms_score_threshold、nms_iou_threshold、max_num_detections、nms_version（v1/v2/batched）、use_cpu_nms；
• Losses 与 PointPillarsTask（L167-L194）：use_wod_metrics 开关 Waymo 官方评测器，init_checkpoint_modules 支持只加载 backbone 或 decoder 做迁移（任务类 initialize 中实现，tasks/pointpillars.py）。

6. 训练：TPU 与 GPU 两种部署

6.1 Cloud TPU 训练

按 official/README-TPU.md 与 GCP 文档完成 TPU 设置后：

MODEL_DIR="gs://<path/to/directory>"
TRAIN_DATA="gs://<path/to/train-data>"
EVAL_DATA="gs://<path/to/eval-data>"

python3 train.py \
--experiment="pointpillars_baseline" \
--mode="train" \
--model_dir=${MODEL_DIR} \
--config_file="configs/vehicle/pointpillars_3d_baseline_tpu.yaml" \
--params_override="task.train_data.input_path=${TRAIN_DATA},task.validation_data.input_path=${EVAL_DATA}" \
--tpu=${TPU}

6.2 多 GPU 训练

python3 train.py \
--experiment="pointpillars_baseline" \
--mode="train_and_eval" \
--model_dir=${MODEL_DIR} \
--config_file="configs/vehicle/pointpillars_3d_baseline_gpu.yaml" \
--params_override="task.train_data.input_path=${TRAIN_DATA},task.validation_data.input_path=${EVAL_DATA}"

README 特别提示：GPU 配置按 8 卡调过参，若使用其他卡数，需要相应调整 batch size、学习率与训练步数。仓库中还有第三份 configs/vehicle/pointpillars_3d_baseline_local.yaml，适合本地小规模验证（对应实验注册函数里 train_steps=100 那种小规模的默认值也是为这类冒烟测试准备的）。

6.3 基线配置逐项解读

以 configs/vehicle/pointpillars_3d_baseline_gpu.yaml 为例（TPU 版 pointpillars_3d_baseline_tpu.yaml 结构相同，仅规模不同）：

runtime:
  distribution_strategy: 'mirrored'   # GPU 用 mirrored，TPU 版为 'tpu'
  mixed_precision_dtype: 'float32'    # 纯 fp32 基线
task:
  model:
    classes: 'vehicle'                # 二分类：车辆/背景
    num_classes: 2                    # 非 'all' 模式必须为 2（task 有断言）
    image:
      x_range: [-76.8, 76.8]          # BEV 覆盖范围（米）
      y_range: [-76.8, 76.8]
      z_range: [-3.0, 3.0]
      resolution: 0.3                 # 0.3 米/格 → 512x512
    pillars:
      num_pillars: 24000
      num_points_per_pillar: 100
      num_features_per_point: 10
    min_level: 1                      # 检测头从 level 1 开始预测
    max_level: 1
    anchors:
    - length: 15.752693               # 单车锚点（车辆长宽先验）
      width: 6.930973
    anchor_labeler:
      match_threshold: 0.6
      unmatched_threshold: 0.45
    featurizer:
      num_blocks: 1                   # 1 个 1x1 卷积块
      num_channels: 64                # BEV 图像通道数 C
    backbone:
      min_level: 1
      max_level: 3
      num_convs: 6                    # 每级下采样组 6 个卷积
    detection_generator:
      pre_nms_score_threshold: 0.05
      nms_iou_threshold: 0.5
      max_num_detections: 200
  train_data:
    global_batch_size: 16             # 2 每卡 x 8 卡
    dtype: 'float32'
    shuffle_buffer_size: 256
    prefetch_buffer_size: 256
  validation_data:
    global_batch_size: 32             # 4 每卡 x 8 卡
trainer:
  train_steps: 494000                 # (158081/16) * 50 epoch
  validation_steps: 1250              # 39987 / 32
  validation_interval: 9880           # 每 epoch 一评
  steps_per_loop: 9880
  summary_interval: 9880
  checkpoint_interval: 9880
  optimizer_config:
    optimizer:
      type: 'sgd'
      sgd:
        momentum: 0.9
        global_clipnorm: 10.0         # 梯度全局范数裁剪
    learning_rate:
      type: 'cosine'
      cosine:
        decay_steps: 494000
        initial_learning_rate: 0.0016
    warmup:
      type: 'linear'
      linear:
        warmup_learning_rate: 0.00016
        warmup_steps: 9880             # 1 epoch 线性预热

GPU 与 TPU 两版配置的核心差异在于规模：

项目	GPU 基线	TPU 基线
硬件	8 × V100（mirrored）	TPU-v2 32 核 pod（4x4 data parallel）
训练全局 batch	16（2×8）	64（2×32）
train_steps（50 epoch）	494000	123500（= (158081/64)×50）
文件头注释的耗时	约 4 hrs/epoch	约 16 mins/epoch，15 hrs/50 epochs
注释记录的精度	mAP 0.46 / mAPH 0.45	mAP 0.45 / mAPH 0.44

注意 train_steps 的推导：训练集 158081 个样本、验证集 39987 个（见配置注释），steps_per_loop = 样本数 / batch 即一个 epoch 的步数，train_steps = epoch 数 × 每 epoch 步数。修改 batch size 后必须同步修改这些步数参数，否则学习率 cosine 衰减周期（decay_steps）与 warmup 长度会失配。

7. 基准结果与实验设置

README "Results" 一节给出的官方 benchmark 设置（训练 TPU 版配置对应此设置）：

• Lidar 范围：X [-76.8, 76.8]，Y [-76.8, 76.8]，Z [-3.0, 3.0]
• Pillars：每帧 24000 根，每根 100 个点，每点 10 个特征
• BEV 图像分辨率：[512, 512, 64]（与 ImageConfig 自动计算结果一致：153.6 / 0.3 = 512）
• 硬件：Cloud TPU-v2（16 核），batch size 64，75 epochs

模型	mAP	mAPH
PointPillars-vehicle	45.96%	45.35%

（README 中附有一个 TensorBoard 实验链接，此处按仓库内容省略外链。）这些数值是仓库文档声明的实测结果，复现时需使用相同的 WOD 1.2.0 版本与预处理参数。

8. 模型导出与推理

• 训练完成后 train.py 会自动调用 utils/model_exporter.py 的 export_inference_graph（batch=1），把 SavedModel 写到 <model_dir>/saved_model，输入为 pillars 与 indices，输出含 NMS 后的检测框、分数、类别与属性；
• 单独导出可使用 tools/export_model.py；
• 推理时注意 Featurizer 的三态设计：call 中 training=None 即 test 模式，固定按 batch=1 处理（modeling/featurizers.py），与导出路径一致。

9. 小结：改动模型时该动哪里

• 换数据集 / 改覆盖范围：先改 task.model.image 与 task.model.pillars（预处理与训练必须一致），再按副本数重新计算 global_batch_size、train_steps、decay_steps、warmup_steps 等步数参数；
• 调检测头输出：在 YAML 中扩展 head.attribute_heads 或 anchors（多尺寸锚点即多尺度检测），对应实现分别在 configs/pointpillars.py 与 modeling/heads.py；
• 调优化策略：trainer.optimizer_config 是 Model Garden 通用结构（official/modeling/optimization/ 下有 SGD/AdamW/cosine/linear warmup 等实现），GPU/TPU 基线都采用 SGD + momentum 0.9 + 全局梯度裁剪 10.0 + 线性预热 + cosine 衰减的组合；
• 验证改动：每个模块都带 *_test.py（如 modeling/featurizers_test.py、tasks/pointpillars_test.py），可用 pointpillars_3d_baseline_local.yaml 做小规模冒烟训练后跑测试确认形状与数值正确。

本实现遵循 Apache License 2.0（见 LICENSE），引用该工作时 README 建议引用原论文：

@inproceedings{alex2019pointpillars,
  title={PointPillars: Fast Encoders for Object Detection from Point Clouds},
  author={Alex H. Lang, Sourabh Vora, Holger Caesar, Lubing Zhou, Jiong Yang, Oscar Beijbom},
  journal={arXiv preprint arXiv:1812.05784},
  year={2019},
}