干货！ASSANet让PointNet更快更强

　　ASSANet: An Anisotropic Separable Set Abstraction for Efficient Point Cloud Representation Learning
　　这篇报告主要介绍NeurIPS 2021 Spotlight论文ASSANet。在ASSANet中，我们深入研究了 PointNet++ [1] 这个经典的网络结构，并开发出更快、更准确的模型变体。
　　我们首先提出了一种新颖的可分离集抽象 (Separable Set Abstraction) 模块，它将 PointNet++ 中使用的普通 SA 模块分解为两个独立的学习阶段：(1) 学习通道相关性和 (2) 学习空间相关性。Separable SA 模块以比普通版本快得多的速度，达到了相当的精度。
　　然后，我们在我们的可分离 SA 模块中引入了一个新的各向异性操作， 并提出了一个各向异性可分离集抽象模块(ASSA)，该模块大大提高了网络的精度。 我们将 PointNet++ 中的SA模块替换为 ASSA 模块，并将修改后的网络命名为 ASSANet。
　　对点云分类、语义分割和部分分割的大量实验表明， ASSANet 优于 PointNet++  和其他方法， 实现了更高的精度和更快的速度 。
　　例如，ASSANet 在 S3DIS Area 5 上的性能比 PointNet++ 高 7.4 mIoU，同时在单个 NVIDIA 2080Ti GPU 上保持 1.6 倍的推理速度。 加深后的 ASSANet 达到 66.8 mIoU，性能优于 KPConv，同时速度提高了 54 倍以上。
　　我们提出了一种 基于各向异性的、可分离且高效的点云特征学习网络 ASSANet。本次研究的动机是，高效的点云处理网络需求越来越大：嵌入式激光雷达传感器的新型移动设备创造了对高效点云学习算法的需求。
　　但是当前 SOTA的方法往往是牺牲速度来换取准确率 。而我们的方法在达到与之前知名方法相同准确率的时候，具备更快的推理速度，甚至比PointNet++快了接近两倍。
　　本期AI TIME PhD直播间，我们邀请到阿卜杜拉国王科技大学(KAUST) 计算机科学专业的博士生——钱国成，为我们带来报告分享《ASSANet：让PointNet++更快更强》。
　　钱国成：
　　阿卜杜拉国王科技大学(KAUST) 计算机科学专业的在读博士生。导师是Bernard Ghanem。钱国成于 2018 年获得西安交通大学的学士学位。研究方向：点云理解，高效神经网络设计。
　　1
　　Methodology
　　ASSANet是基于PointNet++设计的，之所以选择PointNet++，因为这是一种经典的网络结构且被应用到如点云分割等多个领域。其本身性能和速度具有较好平衡。我们本次工作的目的便是研究如何使得PointNet++变得更快更强。
　　Background
　　PointNet++是由SA模块组合起来的，SA指的是Set abstraction，即将点云当作一个集合set，然后将点云的这些点进行下采样以及特征学习。
　　这个模块主要由三个部分组成：subsampling下采样,grouping找邻居,local aggregation 做卷积并聚合(MLPs and reduction)。
　　本文中，我们将PointNet++中的SA模块称为vanilla SA. 值得注意的是， PointNet++原文中提出了多种变体结构，我们的工作是基于PointNet++ multi-grouping，本文将其简称为PointNet++.
　　Latency Analysis
　　研究PointNet++，我们需要首先对其速度进行分析。根据上述对其组成模块的理解，我们发现其latency速度主要受到三个因素影响。
　　point subsampling点云的下采样
　　grouping (neighborhood querying)邻居的query
　　actual computation (mainly MLPs) MLP的计算
　　我们发现MLP的计算依然是限制速度的瓶颈（特别是当点云规模较小的时候），于是我们开始研究对PointNet++ MLP部分进行优化。
　　首先我们分析，
　　(1)为什么vanilla SA速度较慢？这是因为MLP是作用在邻居特征上，但是邻居特征比点云特征多了k倍，说明其中是有相当多的重复数据计算的。若想缓解这种速度问题，最简单的方法就是直接在Point feature X上进行MLP，而不是在邻居feature上做重复计算。
　　(2)为什么PointNet++的准确率不高？因为PointNet++的所有计算都是各向同性的，这包括MLP对每个邻居的feature计算用相同的MLP，以及聚合操作（reduction layer）对每一个邻居的计算相同，这会限制网络的表征能力。
　　根据上述分析，我们逐渐对PointNet++中的Vanilla SA模块进行改进：
　　（1）我们提出了PreConv SA: 直接对点云的特征而不是邻居的特征进行MLP计算。具体表现为 将Vanilla SA中的MLP模块前移至subsampling/grouping之前。这样，模型的FLOPs可以减少k倍。
　　（2）我们进一步提出Separable SA：把MLP分成两部分，一部分是在grouping之前去做，另一部分是在grouping之后去做。需要注意的是，这两部分MLP都是在Point feature上计算，而不是在邻居feature上计算，因此大大减少了计算量。然后，我们在这两部分MLP的输出上增加了一个残差连接。随着网络深度的增大，残差连接可以提升网络精度。
　　（3）我们提出了各向异性的聚合模块(ASSA)。PreConv SA, Separable SA只是让网络变得更快了，为了进一步提升网络的精度，我们需要解决之前提到的vanilla SA的各向同性问题。
　　具体做法如下：在对邻居特征直接进行各向同性的reduction之前，我们将邻居特征乘以领居与中心点的在x, y, z方向的相对距离，从而实现对邻居特征不同程度的缩放。
　　我们分别用x, y, z三个方向的相对距离缩放邻居特征，并将三个缩放的邻居特征级联在一起，并最后对缩放后的邻居特征进行avgpooling。
　　值得注意的是，我们提出的各向异性聚合模块是基于PosPool [2]提出的，区别在于（1）我们没有对特征进行分组缩放；（2）我们用ball query的半径归一化了缩放比例。实验证明，我们提出的ASSA聚合模块性能优于PosPool。
　　在PointNet++中，我们把vanilla SA替换成我们最终提出的ASSA。结果如下图所示，在有4096个点输入的时候，我们的速度会提升4倍；在有15000个点输入的时候，速度可以提升1.6倍。
　　Network: ASSANet & ASSANet (L)
　　ASSANet：我们将PointNet++中的普通SA模块替换为我们的ASSA模块，并新增stem aggregation层（相当于图像识别领域常用的stem layer），即在点云下采样之前多增加一层local aggregation。修改后的网络结构被称为ASSANet。
　　ASSANet (L)：我们考虑了两种模型扩展策略，一种是加宽（增大初始特征数量C）；另一种是加深（增大每一个SA module中聚合层的数量D）。我们将扩展后的ASSANet (C=128, D=3)称为 ASSANet (L).
　　2
　　Experiment
　　ASSANet在多个benchmark上都比PointNet++速度更快精度更高，并同时取得接近SOTA的结果。Ablation study也揭示了我们提出的ASSA模块的优越性。
　　Feature pattern visualization
　　为了帮助理解ASSA的作用，我们可视化了特征。如上图所示，在做ASSA module之前只有部分区域被激活。而在做了一层ASSA之后，图中的飞机机翼和尾翼，台灯的灯泡和底座，都被同时捕捉到了。这说明了ASSA模块能够帮助模型捕捉几何信息。
　　参考文献：
　　[1] Qi, C. R., Yi, L., Su, H., & Guibas, L. J. (2017). Pointnet++: Deep hierarchical feature learning on point sets in a metric space. Advances in neural information processing systems, 30.
　　[2] Liu, Ze, Han Hu, Yue Cao, Zheng Zhang, and Xin Tong. ＂A closer look at local aggregation operators in point cloud analysis.＂ In European Conference on Computer Vision, pp. 326-342. Springer, Cham, 2020.
　　论文链接：
　　https://arxiv.org/abs/2110.10538
　　代码开源：
　　https://github.com/guochengqian/ASSANet