文 / 研究员 Alireza Fathi 和 AI 实习生 Rui Huang,Google Research
过去几年里,3D 传感器(如激光雷达、深度传感摄像头和雷达)越发普及,对于能够处理这些设备所捕获数据的场景理解技术,相应需求也在不断增加。此类技术可以让使用这些传感器的机器学习 (ML) 系统,如无人驾驶汽车和机器人,在现实世界中导航和运作,并在移动设备上创建改进的增强现实体验。
最近,计算机视觉领域开始在 3D 场景理解方面取得良好进展,包括用于移动 3D 目标检测、透明目标检测等的模型,但由于可应用于 3D 数据的可用工具和资源有限,进入该领域本身可能具有挑战性。
为了进一步提高 3D 场景理解能力,降低感兴趣的研究人员的入门门槛,我们发布了 TensorFlow 3D (TF 3D),这一高度模块化的高效库旨在将 3D 深度学习功能引入 TensorFlow。TF 3D 提供了一组流行的运算、损失函数、数据处理工具、模型和指标,使更广泛的研究社区能够开发、训练和部署最先进的 3D 场景理解模型。
TF 3D 包含用于最先进 3D 语义分割、3D 目标检测和 3D 实例分割的训练和评估流水线,并支持分布式训练。它也可实现其他潜在应用,如 3D 目标形状预测、点云配准和点云密化。此外,它还提供了统一的数据集规范和配置,用于训练和评估标准 3D 场景理解数据集。目前支持 Waymo Open、ScanNet 和 Rio 数据集。不过,用户可以将 NuScenes 和 Kitti 等其他流行数据集自由转换为相似格式,并将其用于预先存在或自定义创建的流水线,也可以通过利用 TF 3D 进行各种 3D 深度学习研究和应用,包括快速原型设计以及试验新想法的方式来部署实时推断系统。
左侧为 TF 3D 中 3D 目标检测模型在来自 Waymo Open 数据集的一帧画面上的示例输出。右侧为 3D 实例分割模型在来自 ScanNet 数据集 的场景上的示例输出
我们将介绍 TF 3D 提供的高效可配置的稀疏卷积主干,它是在各种 3D 场景理解任务上取得最先进结果的关键。此外,我们将分别介绍 TF 3D 目前支持的三种流水线:3D 语义分割、3D 目标检测和 3D 实例分割。
传感器捕获的 3D 数据通常具有一个场景,其中包含一组感兴趣的目标(如汽车、行人等),其周围大多是有限(或无)兴趣的开放空间。因此,3D 数据本质上是稀疏的。在这样的环境下,卷积的标准实现将需要大量计算并消耗大量内存。因此,在 TF 3D 中,我们使用子流形稀疏卷积和池化运算,旨在更有效地处理 3D 稀疏数据。稀疏卷积模型是大多数户外自动驾驶(如 Waymo、NuScenes)和室内基准(如 ScanNet)中应用的最先进方法的核心。
子流形稀疏卷积
https://arxiv.org/abs/1706.01307
我们还使用多种 CUDA 技术来加快计算速度(例如,哈希处理、在共享内存中分区/缓存过滤器,以及使用位运算)。Waymo Open 数据集上的实验表明,该实现比使用预先存在的 TensorFlow 运算的精心设计实现快约 20 倍。
TF 3D 然后使用 3D 子流形稀疏 U-Net 架构为每个体素 (Voxel) 提取特征。通过让网络同时提取粗略特征和精细特征并将其组合以进行预测,事实证明 U-Net 架构是有效的。U-Net 网络包括编码器、瓶颈和解码器三个模块,每个模块都由许多稀疏卷积块组成,并可能进行池化或解池化运算。
3D 稀疏体素 U-Net 架构。请注意,水平箭头接收体素特征并对其应用子流形稀疏卷积。下移箭头执行子流形稀疏池化。上移箭头将回收池化的特征,与来自水平箭头的特征合并,并对合并后的特征进行子流形稀疏卷积
上述稀疏卷积网络是 TF 3D 中提供的 3D 场景理解流水线的主干。下面描述的每个模型都使用此主干网络提取稀疏体素的特征,然后添加一个或多个附加预测头来推断感兴趣的任务。用户可以更改编码器/解码器层数和每层中卷积的数量以及修改卷积过滤器的大小来配置 U-Net 网络,从而通过不同的主干配置探索大范围的速度/准确率权衡。
3D 语义分割模型只有一个输出头,用于预测每个体素的语义分数,语义分数映射回点以预测每个点的语义标签。
来自 ScanNet 数据集的室内场景的 3D 语义分割
3D 语义分割模型
https://arxiv.org/abs/1711.10275
在 3D 实例分割中,除了预测语义外,目标是将属于同一目标的体素归于一组。TF 3D 中使用的 3D 实例分割算法基于我们先前的使用深度指标学习的 2D 图像分割研究工作。该模型预测每个体素的实例嵌入向量以及每个体素的语义分数。实例嵌入向量将体素映射到一个嵌入向量空间,其中对应同一目标实例的体素靠得很近,而对应不同目标的体素则相距很远。在这种情况下,输入是点云而不是图像,并使用 3D 稀疏网络而不是 2D 图像网络。在推断时,贪婪的算法每次挑选一个实例种子,并使用体素嵌入向量之间的距离将其分组为段。
深度指标学习的 2D 图像分割研究
https://arxiv.org/abs/1703.10277
3D 目标检测模型预测每个体素的大小、中心、旋转矩阵以及目标语义分数。在推断时,采用盒建议机制 (Box proposal mechanism) 将成千上万的各体素的盒预测减少为几个准确的盒建议,然后在训练时,将盒预测和分类损失应用于各体素的预测。我们对预测和基本事实盒顶角之间的距离应用 Huber 损失。由于从其大小、中心和旋转矩阵估计盒顶角的函数是可微的,因此损失将自动传播回这些预测的目标属性。我们采用动态盒分类损失,将与基本事实强烈重合的盒分类为正,将不重合的盒分类为负。
我们在 ScanNet 数据集上的 3D 目标检测结果
在我们最近的论文《DOPS: Learning to Detect 3D Objects and Predict their 3D Shapes》中,我们详细描述了 TF 3D 中用于目标检测的单阶段弱监督学习算法。此外,在后续工作中,我们提出基于 LSTM 的稀疏多帧模型,扩展了 3D 目标检测模型以利用时间信息。我们进一步证明,在 Waymo Open 数据集中,这种时间模型比逐帧方法的性能高出 7.5%。
DOPS 论文中介绍的 3D 目标检测和形状预测模型。3D 稀疏 U-Net 用于提取每个体素的特征向量。目标检测模块使用这些特征建议 3D 盒和语义分数。同时,网络的另一个分支预测形状嵌入向量,用于输出每个目标的网格
DOPS:Learning to Detect 3D Objects and Predict their 3D Shapes
https://arxiv.org/abs/2004.01170
这个代码库在我们的 3D 计算机视觉项目中发挥了明显的作用,我们希望对您也一样。我们欢迎对代码库的贡献,也请继续关注我们对框架的进一步更新。要开始使用,请访问我们的 GitHub 仓库。
GitHub 仓库
https://github.com/google-research/google-research/tree/master/tf3d
TensorFlow 3D 代码库和模型的发布是 Google 研究人员在产品组的反馈和测试下广泛合作的结果。我们要特别强调 Alireza Fathi 和 Rui Huang(在 Google 期间完成的工作)的核心贡献,另外还要特别感谢 Guangda Lai、Abhijit Kundu、Pei Sun、Thomas Funkhouser、David Ross、Caroline Pantofaru、Johanna Wald、Angela Dai 和 Matthias Niessner。
想深入了解 TensorFlow 3D 实操技术,请访问 TensorFlow 官方网站。想了解 TensorFlow 最新资讯,请扫描下方二维码,关注 TensorFlow 官方公众号。