Google AI推出了TensoRFlow 3D,将3D深度学习功能整合进TensoRFlow,增加了3D稀疏卷积网络。根据在WayMo Open数据集上的实验,这一实现的速度提升达到了20倍。
随着自动驾驶和机器人技术的不断进步,激光雷达、深度摄像头和雷达等3D传感器已成为获取道路信息的重要工具。
这些传感器驱动的机器学习系统在现实世界中的导航等关键功能显得尤为重要。
最近,在目标检测和透明目标检测等方面的3D场景理解取得了显著进展,尽管由于缺乏可用的工具和资源,该领域仍面临不少挑战。
TensoRFlow 3D:将TensoRFlow与3D深度学习结合
为了提升3D场景建模的效率,并简化研究者的工作流程,Google AI推出了TensoRFlow 3D (TF 3D),这是一个高效且高度模块化的库,旨在为TensoRFlow引入3D深度学习能力。
TF 3D提供了一系列当前常用的操作、损失函数、数据处理工具、模型和评估指标,使得更多研究团队能够开发、训练和部署最先进的3D场景理解模型。
TF 3D支持最前沿的3D语义分割、3D目标检测和3D实例分割的训练与评估任务,并且支持分布式训练。
此外,TF 3D还适用于其他潜在应用,如三维物体形状预测、点云配准和点云增密,同时提供一致的数据集规范及三维场景理解数据集的训练与评估标准配置。
目前,TF 3D支持WayMo Open、ScanNet和Rio数据集。
用户还可以将其他流行数据集,如NUScenes和KITTI,转换为类似格式以供使用,亦可利用TF 3D进行各种3D深度学习研究和应用,从快速原型设计到实时推理系统的部署。

左侧展示的是TF 3D中3D物体检测模型在WayMo Open数据集的一帧输出示例,右侧为ScanNet数据集上3D实例分割模型的输出示例。
接下来,我们将介绍TF 3D中高效且可配置的稀疏卷积骨干,这是在多种3D场景理解任务中获得最佳结果的关键。
此外,我们还将逐一介绍TF 3D当前支持的三项流水线任务:3D语义分割、3D目标检测和3D实例分割。
3D稀疏卷积网络
传感器采集的3D数据通常呈现出一个场景,其中包含一组感兴趣的物体(如汽车、行人等),而周围则多为空旷空间。因此,3D数据本质上是稀疏的。
在这种环境中,标准卷积实现需要大量计算和内存消耗。因此,在TF 3D中,我们采用了流形稀疏卷积(subManifold sparse convolution)和池化操作,这些方法能更高效地处理3D稀疏数据。
稀疏卷积模型是大多数户外自动驾驶(如WayMo、NUScenes)及室内基准测试(如ScanNet)中应用的最新技术的核心。

谷歌还运用了多种CUDA技术来加速计算速度,如hash、共享内存中的分区/缓存过滤器以及位操作等。
实验结果显示,在WayMo Open数据集上,这种实现的速度比预设的TensoRFlow操作快约20倍。

图源:WayMo Open Dataset on GitHub
接下来,TF 3D利用3D流形稀疏U-Net架构提取每个voxel的特征。U-Net架构通过提取粗细特征并将其组合进行预测,已被证明是有效的。
U-Net网络由编码器、瓶颈和解码器三个模块构成,每个模块包含多个稀疏卷积块,并可能进行池化或非池化操作。

上述3D稀疏体素U-Net架构中,水平箭头接收体素特征,并应用流形稀疏卷积。向下的箭头执行流形稀疏池化,而向上的箭头则将池化后的特征与水平箭头的特征进行连接,并对连接后的特征进行流形稀疏卷积。
这一稀疏卷积网络构成了TF 3D提供的3D场景理解流水线模型的骨干。
接下来描述的每个模型均使用该骨干网络提取稀疏体素特征,然后添加一个或多个预测头来推断感兴趣的任务。
用户可以通过调整编码器/解码器层数和每层卷积的数量来配置U-Net网络,并通过修改卷积滤波器的尺寸,以在速度和精度之间进行权衡。
三维语义分割
三维语义分割模型仅有一个输出,用于预测每个点的语义分数,并将其映射回点,以预测每个点的语义标签。

从ScanNet数据集中对室内场景进行的3D语义分割示例。
三维实例分割
在三维实例分割中,除了预测语义外,更重要的是将同一对象的体素组合在一起。
TF 3D中使用的3D实例分割算法基于深度度量学习方法的2D图像分割工作。该模型能够预测每个体素的实例嵌入向量及其语义评分。
实例嵌入向量将体素映射到一个嵌入空间,使得同一对象实例的体素距离较近,而不同对象的体素距离较远。
在此情况下,输入为点云而非图像,并使用三维稀疏网络而不是二维图像网络。在推理过程中,利用贪心算法选择实例种子,并通过体素嵌入的距离函数将不同的体素聚合到相应的实例上。
三维目标检测
目标检测模型能够预测每个体素的大小、中心和旋转矩阵,以及对象的语义评分。
在推理时,推选机制将多个候选框处理为少数几个精确的3D目标框。在训练过程中,使用预测与GT之间的Huber损失来计算损失。由于使用大小、中心和旋转矩阵来估算框的边角是可微的过程,损失可以自然地在预测过程的权重中传递。研究人员还利用动态框分类损失来区分预测的正例和负例。

ScanNet数据集上3D物体检测的结果。
TF 3D只是市场上众多3D深度学习扩展中的一个。2020年,Facebook推出了PyToRch3D,专注于3D渲染和虚拟现实。而英伟达的Kaolin则是一个模块化的高分辨率模拟环境的应用。
从整体来看,TF 3D的应用似乎更加聚焦于机器人感知和映射,而其他选项则更偏向于3D模拟与渲染。为了实现3D渲染,Google推出了TensoRFlow Graphics。
