让多任务奔跑起来！Waymo开放数据集3D语义分割任务冠军方案分享（图森）

论文链接：https://arxiv.org/pdf/2206.11428.pdf

汽车人的碎碎念

今天分享图森打榜Waymo开放数据集3D语义分割任务的技术报告，整篇文章读下来比较通透。比较经典的打榜思路，尤其是后面的多任务学习、两阶段优化，有打榜或训练大模型需求的小伙伴可以仔细读读~

摘要

本文是2022年Waymo开放数据集3D语义分割挑战的第一名所形成的技术报告。算法名称LidarMultiNet，统一输出了3D语义分割、目标检测和全景分割任务。LidarMultiNet的核心是一个强大的基于3D体素的encoder-decoder网络，包含一个新颖的Global Context Pooling（GCP）模块，可从LiDAR帧中提取全局上下文特征以补充其局部特征。进一步，提出了一个可选择的二阶段模块来细化一阶段的分割结果。最终，LidarMultiNet取得了3D语义分割挑战第一名的成绩，并首次证明了主要的LiDAR感知任务可以统一在一个端到端的网络中实现。

本文的贡献如下：

1. 提出了一个高效的基于体素的LiDAR多任务网络，其统一了主要的 LiDAR 感知任务；
2. 提出了一种新颖的Global Context Pooling (GCP) 模块，以改进基于3D稀疏卷积的编码器-解码器架构中的全局特征学习；
3. 引入了两阶段的细化模块来细化第一阶段的语义分割结果或输出更准确的全景分割结果。

方法

Voxel-based LiDAR Segmentation

LiDAR点云语义分割旨在预测点云中所有点云的语义标签，每个点云有个输入特征，即3D坐标、反射强度、时间戳等等。

体素化：首先将点云体素化，之后使用一个简单的体素特征编码器，其只包含多层感知器（MLP）和最大池化层来生成稀疏的体素特征表示：

基于稀疏体素的编码器-解码器主干网络：本文设计了2D UNet [22]的3D形式作为分割网络的主干，其中体素特征在编码器中逐渐下采样至1/8大小，并在解码器中逐步恢复。Key indices在编解码器间以相同的比例共享，以保持相同的稀疏性并减少稀疏索引的重复计算。在训练期间，LidarMultiNet使用体素级的标签训练模型，在推理时通过去体素化的步骤投影回点云上。

LidarMultiNet

LidarMultiNet的网络结构如下图所示。主网络由基于3D稀疏卷积和反卷积的编码器-解码器王所组成，提出的GCP模块作用于编码器和解码器之间。网络可以增加两个辅助任务，即3D目标检测和BEV分割任务。此外，还有可选的第二个阶段用来细化语义分割结果并输出全景结果。

Global Context Pooling（GCP）

3D稀疏卷积大大减少了3D CNN 的内存消耗和运行时间，但其通常需要相同尺度的层才能在编码器和解码器中保持相同的稀疏性。这就将网络限制成只能使用相同尺度的submanifold convolution[10]，如下图所示。然而submanifold convolution不能通过堆叠多个卷积层将特征发散到孤立的体素上，所以在一定程度上限制了网络学习全局上下文的能力。Global Context Pooling（GCP）旨在通过将全局上下文特征投影到2D密集的BEV特征图上，接着在应用2D多尺度特征提取器后将密集特征图转换回来，最后从3D稀疏张量中提取全局上下文信息。

如下图所示，基于编码器输出的3D稀疏张量，首先将其投影到密集的BEV特征图中。进一步在高度维度上将其与BEV特征图拼接在一起，形成新的2D BEV特征图。然后，使用2D多尺度CNN进一步提取上下文信息。最后将编码后的BEV特征重塑为密集体素特征，接着再使用密集到稀疏的转换步骤将其转换为稀疏体素特征。

得益于GCP，LidarMultiNet可以显著扩大感受野，而感受野对语义分割至关重要。此外，GCP中的BEV特征图可以与目标检测等任务恭喜，只需要增加额外的分支即可，稍微增加一些计算量。通过增加目标检测等任务，GCP可以进一步提高分割性能。

LiDAR Multi-task Learning

辅助任务：本文首先使用anchor-free算法CenterPoint [31]的检测头作为额外多任务分支。动机主要有两点：首先，检测损失可以作为辅助损失，并有助于利用分割和检测任务之间的协同作用，Waymo开放数据集 (WOD)也提供了对应的真值；其次，我们的算法超越了同时执行两项任务的多任务网络。除了检测头外，还可以再增加一个BEV分割头，输出粗糙的分割结果，并在训练期间作为另一个辅助损失。

多任务训练和损失函数：BEV的分支使用CE和Lovasz损失训练。整体的网络训练是端到端的，包含分割损失、检测损失和BEV损失，并使用如下的带权重的多任务损失系数进行平衡：

Second-stage Refinement

在训练期间，检测作为分割网络的辅助损失，不直接对分割做贡献，如上图所示，由于缺乏空间先验知识，检测框中的点有可能被误分类为多个类别。为了提高实例类别的空间一致性，本文提出了一种的新的基于点云的方法作为第二个阶段，该方法也可输出全景结果，如下图所示。

实验结果

实现细节如下表所示：

消融实验如下，各个模块有比较明显的提点。

打榜当然要使用TTA（Test-Time Augmentation）和模型集成咯。

最终取得了Waymo开放数据集3D语义分割任务榜单第一的成绩。

参考

[1] LidarMultiNet: Unifying LiDAR Semantic Segmentation, 3D Object Detection, and Panoptic Segmentation in a Single Multi-task Network

原文首发微信公众号【自动驾驶之心】：一个专注自动驾驶与AI的社区（https://mp.weixin.qq.com/s/NK-0tfm_5KxmOfFHpK5mBA）