平均7倍实测加速，MIT提出高效、硬件友好的三维深度学习方法

机器之心发布

机器之心编辑部

随着三维深度学习越来越成为近期研究的热点，基于栅格化的数据处理方法也越来越受欢迎。但这种处理方法往往受限于高分辨下巨大的内存和计算开销，因此麻省理工学院 HAN Lab 的研究者提出利用 Point-Voxel CNN 来实现高效的三维深度学习，同时能够避免巨大的数据访问开销并很好地提升了局部性。该论文已被 NeurIPS 2019 接收为 Spotlight Presentation。

• 网站：http://pvcnn.mit.edu
• 论文: https://arxiv.org/pdf/1907.03739.pdf

摘要

随着传感器技术的发展和大量新兴应用场景（AR/VR/自动驾驶）的出现，三维深度学习成为了近期的研究热点。三维数据往往以点云的方式存储，近年来，研究人员抑或是选择先将点云离散化成结构化的、规整的栅格形式（voxels，可以类比 2D 的像素 pixels），再利用体素卷积神经网络（volumetric CNNs，可以看作 2D CNN 的三维推广）对栅格数据进行处理；抑或是选择直接在点云数据上进行卷积计算。

基于栅格化数据的处理方法往往受限于其在高分辨率下巨大的内存和计算开销：从复杂度的意义上说，volumetric CNNs 的运算量和内存占用都和栅格的分辨率的三次方成正比，导致 volumetric CNNs 无法和 2D CNN 一样处理高分辨率的输入。而对于直接在点云上进行深度学习的方法，尽管点云表示大大减少了存储空间上的占用，但要在稀疏的、非结构化的点云数据上进行卷积却绝非易事。

我们发现之前的工作竟需要花 80% 以上的时间进行对数据的不规则访问以作为卷积运算的准备工作，而实际计算所占的时间比例却非常低，这无疑造成了基于点云的深度学习方法往往比较低效。为解决这样的问题，我们提出了 Point-Voxel CNN（PVCNN）来实现高效的三维深度学习。PVCNN 利用点云的形式来存储数据，以减小内存的消耗，而又选择栅格化的形式进行卷积，这样可以避免处理点云稀疏性所带来的巨大的非规则数据访问开销，很好地提升局部性。这种取长补短的设计使我们的 PVCNN 在内存和计算上都十分高效：我们只用十分之一的内存就可以取得远高于 volumetric CNN baseline 的准确率，而对比于多个基于点云直接进行深度学习的方法，我们又可以得到平均 7 倍的实测加速。

PVCNN 同时展现了在嵌入式、端计算设备上部署的潜力：通过缩小 PVCNN 的通道数得到的小模型往往可以在三维语义分割任务上比 PointNet（一个极其高效的模型）快两倍，却达到它性能更加强大的续作 PointNet++的精确度。我们也在与自动驾驶密切相关的三维物体检测任务中验证了我们方法的有效性。基于我们 Frustum-PVCNN 的模型平均相对 Frustum-PointNet++有 2.4% 的 mAP 提升，同时还有 1.5 倍的加速与内存节省。相关工作已经被 Neural Information Processing Systems 2019 接收为 Spotlight Presentation，我们将在 12 月 10 日于温哥华宣讲我们的工作。

背景

三维深度学习在近期研究中取得很大进展。然而，之前的工作更关注如何提升三维深度学习计算子的表征能力，却鲜有研究致力于提升三维深度学习方法的效率。

从计算机硬件的角度看，对于存储器的访问往往对算法的效率有着决定性的影响。下面的左图显示，在硬件上，片下的 DRAM 访存能量消耗比片上的 SRAM 访问或者一次浮点乘加计算贵两个数量级，而 DRAM 访问的带宽也远小于片上资源或是运算资源（这导致了访存可能带来高延时）；另一方面，下面的右图揭示了在 DDR 内存中非规则访存带来的危害：频繁的板块冲突（bank conflict）导致数据无法在一个时钟周期内被连续从不同内存板块（bank）中读出，从而进一步提升访存带来的延时。

因此，为使得三维深度学习方法更加高效，我们需要避免大量访存和不规则访存。

对于 2016 年前后被大量研究的基于栅格的三维深度学习方法，制约其效率的掣肘是随分辨率三次方增加的内存/计算渐进复杂度。栅格作为一种类似二维像素的离散化表征形式，其不可避免的困难是落在同一个格中的两个点将在栅格化后变得无法区分，从而导致信息丢失。我们在 ShapeNet 数据集上做了统计，发现如果要在一块 Titan Xp (12GB 显存) 上训练体素卷积神经网络，只能容纳分辨率为 64x64x64 的栅格输入，此时会有 42% 的信息损失（即只有 58% 的点是可区分的），而增加七倍的计算资源也只是能让可区分的点数勉强超过 90%。如前所述，这样的内存消耗是由栅格表征形式本身的复杂度所决定的。

基于点云的深度学习方法则不同，影响其效率的实际上是不规则访存。我们知道，点云上的每个点在三维空间里连续分布，而在内存上相邻的点在真正的三维空间里并不一定相邻。这就给卷积运算，这种定义在每个点邻域上的运算带来挑战。一方面，获得每个点的邻域需要大量的非规则数据访问；另一方面，由于三维空间的连续性，我们无法简单地通过类似 2D CNN 或 Volumetric CNN 的方式去定义一个基于查找表（Look-up Table）的卷积核，只能去动态定义卷积滤波器，这又会带来一些计算开销。

下图显示了我们的发现：之前最先进的基于点云的深度学习方法 DGCNN, PointCNN 和 SpiderCNN，往往需要 50%-80% 的时间进行非规则数据访问和动态卷积核计算，而实际计算的时间甚至可能低至 20% 以下。很显然，从运行时间的角度，基于点云的方法也具有极大的优化空间。

为此我们提出 Point-Voxel Convolution（PVConv）作为一种全新的、高效的三维深度学习计算子。如前所述，PVConv 仍使用点云形式存储数据，因而内存占用小，整体的访存量也小；另一方面，我们的 PVConv 在栅格表征下进行卷积；这样，我们大大省去了传统的基于点云的三维深度学习方法在进行卷积时难以避免的非规则内存访问及动态卷积核生成代价，从而获得上图所示的极高的实际计算比。

通过 2D CNN 般地重复 PVConv，我们可以得到 PVCNN——一个可用于多种三维计算机视觉任务的强大而高效的深度学习网络。我们在多个物体级别、室内场景级别以及室外场景中物体检测的任务中展示了 PVCNN 的有效性。

方法

Point-Voxel Convolution

我们选用两种 Irregular Access Cost 和 Dynamic Kernel Computation Cost 为零的子操作：多层感知机和栅格化卷积。下面我们将展示如何以很低的代价将两者有机结合并得到我们的 Point-Voxel Convolution 这种计算子。上图中上面的分支为基于栅格的信息聚合分支，其分辨率相对较低：

归一化：

输入点云的尺度往往会有一些差异，而我们希望在进行进行栅格化卷积时特征图的尺度是相同的，这就需要进行归一化。在我们的方法中，我们只在栅格分支局部地进行这种归一化操作：将输入点云缩放到单位球中，并通过 [-1, 1] 到 [0, r-1] 的坐标线性映射把点云坐标转成栅格坐标。

栅格化：

我们将点云表示局部地转换成栅格表示以便进行栅格化卷积。直观地，我们利用 Charles Qi 等人提出的对称函数建模同一栅格内的所有点：我们选择将栅格内所有点进行平均池化以得到代表某个栅格的最终特征。

信息整合：

我们在得到的栅格特征图上进行栅格卷积，以达到信息整合的目的。类似于 2D CNN 和普通的 Volumetric CNN，我们在每次栅格卷积后也通过批归一化（BatchNorm）和非线性激活函数（ReLU）。

反栅格化/点云化：

如前所述，栅格表征无可避免地会丢失一部分信息。为了保持点云始终处于高分辨率，我们需要通过反栅格化操作将栅格数据重新转换为点云。一个朴素的想法是直接给落在每个栅格内的点赋上其所在栅格的特征，然而，这样会导致同一栅格内点的特征完全相同。我们也可以选择三线性插值（Trilinear Interpolation）来根据每个点到相邻栅格格点的位置来决定其特征（由相邻的八个栅格点插值得到）。这样，我们可以保证每个点在反栅格化后的特征通常不相同。

上图中的下面分支为基于多层感知机的特征转换分支，它通过一层全连接层保持了点云处于高分辨率，这样的高分辨率信息可以为来自栅格化分支的低分辨率、邻域信息提供重要的补充。当我们得到了两个分支的信息后，我们可以简单地通过 elementwise addition 将它们整合起来，这就是一个完整的 PVConv 操作。

值得注意的是，PVConv 是一个非常高效的计算子。如前述，PVConv 显然得益于两个分支本身不存在非规则访存以及动态卷积核计算。而唯一引入额外开销的栅格化和反栅格化（朴素的）过程，其非连续访存的代价都是 O(n) 次，因为每个点只需要被访问一次；而对于一般的基于点云的三维深度学习计算子，进行一次邻域信息聚合首先要把每个点的邻域信息取来，这里就有 O(kn) 次的非连续内存访问，k 是邻域中的点数，往往在 16 到 64。很显然，PVConv 大大减小了非连续内存访问。另外一方面，PVConv 并不需要像基于点云的三维深度学习计算子那样去计算每个点的邻域，这也大大降低了额外开销。

从宏观上看，我们的 PVConv 中每个步骤都是可导的，并且它的输入和输出都是点云。所以从梯度流的角度看，PVConv 和一般的直接基于点云的三维深度学习计算子一样，既可以通过自身堆叠形成比较深的神经网络，也容易地被整合进任何端到端可训练的三维深度学习网络中（见下两图：第一张是 PVConv 单独组合成 PVCNN，第二张图是 PVConv 嵌入 PointNet++得到 PVCNN++）。

结果

我们在具有代表性的三维视觉数据集：ShapeNet（三维物体部分语义分割），S3DIS（三维室内场景分割）和 KITTI（自动驾驶场景中的三维物体检测）中验证了 PVCNN 的表现。

在 ShapeNet 数据集上，我们的完整通道数模型可以以 2.7 倍的实测加速达到略好于 PointCNN 的性能，同时我们的内存占用也减小了 1.6 倍；而即便我们将通道数缩小四倍，我们的模型也可以在达到 PointNet++, DGCNN 等复杂方法性能的同时比 PointNet（此前这个领域中被认为最为精简、高效的模型）快 1.9 倍、省 1.9 倍内存。

我们同时展示我们的四分之一宽度小模型和我们的完整模型在端计算设备上部署的情况：

完整的 PVCNN 可以在 NVIDIA Jetson AGX Xaviers 上以每秒 20.2 个物体的速度运行，而四分之一宽度版本的小模型在价格仅 99 美元的 NVIDIA Jetson Nano 上也可以达到接近每秒 20 个物体的速度。

通过下图中对 PVConv 两个分支特征的可视化，我们不难发现两个分支确实学到了互补的特征：

基于栅格的分支往往更加关注大而连续的区域，基于点云 MLP 的特征转换分支往往更加关注独立的，不连续的区域。

在 S3DIS 数据集上，我们的方法也具有先进的性能，同时极为高效：只需要七分之一左右的运行时间就可以超越此前优秀的方法 PointCNN。当然，我们的内存消耗也比 PointCNN 小了很多。此外，我们的 0.25 倍通道数小模型均在运行时间比 PointNet 快 1.8 倍的情况下性能超越了更为复杂的 DGCNN 和 RSNet。Tradeoff 曲线更直观地反映了这一点。

在 KITTI 数据集上，我们用 PVConv 替换 Frustum PointNet 系列中三维部分的计算子，并达到明显更优的性能：省 1.5 倍内存和运行时间的条件下，平均检测 mAP 比 Frustum PointNet++高 2.4%。