NeurIPS2021 | 华南理工提出SS-Conv：兼顾加速与SE(3)等变性，3D空间姿态估计突出

2021 年 12 月 5 日 机器之心

机器之心专栏

华南理工大学等

来自华南理工大学等机构的研究者提出了一个新颖的稀疏姿态可控卷积（SS-Conv），SS-Conv 不仅利用稀疏张量对姿态可控卷积进行极大地加速，还在特征学习中严格地保持 SE(3) 等变性。

近年来，SE(3) 等变网络在 3D 语义分析任务中发挥着重要的作用，尤以 3D 姿态可控卷积神经网络为代表。3D 姿态可控卷积神经网络利用姿态可控卷积（ST-Conv）逐层学习姿态等变的特征，从而保留 3D 输入的姿态信息。直观地说，对于一个 ST-Conv，关于 3D 输入的任意 SE(3) 变换(r,t)，都将诱导输出特征的同步姿态变换，其中r ∈SO(3) 表示旋转，

表示平移。图 1(a) 进一步给出了解释，给定输入的一个 SE(3) 变换，特征向量所在位置均关于(r,t)进行刚性变换，而特征向量 ρ(r) 本身也关于进行旋转（ ρ (r) 是旋转r的一个表征）。 SE(3) 等变的特性激活了特征空间的姿态可控性，例如，SE(3) 变换可通过对特征空间进行操控来直接实现，而无需变换输入。为了生成姿态可控特征，ST-Conv 将特征域限定在 3D 体素数据的规则网格上，使其能通过 3D 卷积方便地实现。对 3D 卷积的兼容简化了 ST-Conv 的实现，但也牺牲了对不规则且稀疏的 3D 数据（例如，点云）的高效处理，导致 ST-Conv 未能在更多 3D 语义分析领域中被广泛使用。

图 1. SE(3) 等变性：（a）ST-Conv；（b）SS-Conv。箭头表示 3D 场中有向的特征向量。

为了解决上述问题， 华南理工大学等研究人员提出了一个新颖的稀疏姿态可控卷积（SS-Conv） 。SS-Conv 不仅利用稀疏张量对姿态可控卷积进行极大地加速，还在特征学习中严格地保持 SE(3) 等变性。图 1(b) 对 SS-Conv 的 SE(3) 等变特性进行了阐释。为了实现 SS-Conv，研究人员通过基于球形谐波的基核的线性组合来建立卷积核，使其满足 SE(3) 等变卷积应遵循的旋转可控约束条件，同时基于稀疏张量在激活的特征位置上利用 GPU 上的矩阵加乘操作实现快速卷积。

论文链接：https://arxiv.org/abs/2111.07383
代码链接：https://github.com/Gorilla-Lab-SCUT/SS-Conv

虽然 SE(3) 等变特征学习在 3D 物体识别任务上被广泛使用，其潜能在 3D 语义分析的其他任务上却尚待开发。研究人员尝试将 SS-Conv 应用到 3D 空间中物体姿态估计中，因此提出了一个基于 SS-Conv 的通用框架，通过堆叠多层 SS-Conv 来提取 SE(3) 等变特征，从而直接解码出物体的姿态。在这个框架中，一个新颖的特征操控模块（Feature-Steering module）充分地利用特征的可控性来迭代地对姿态进行优化。研究人员在三个姿态相关的 3D 物体语义分析任务上进行充分的实验，包括实例层级的 6D 物体姿态估计、类别层级的 6D 物体姿态及大小估计、类别层级的 6D 物体姿态跟踪。

SS-Conv 的实现

假设卷积核 k 为一个离散的规则立方体，具有s×s×s 个网格点

，则 SS-Conv 可定义为：

其中，

表示位置x在特征空间

中的状态。

代表位置x未激活，特征

处于基态；

代表位置x已被激活，

处于激活态。

与 ST-Conv 相比，SS-Conv 在两个方面进行加速：i) 卷积只在被激活的输出位置上进行，而不是整个 3D 体素空间中，其中激活位置的数量仅占一小部分；ii）在每个激活的输出位置的感受野中，只有激活的输入特征被卷积。为了上述目的，输入和输出特征被分别表示为稀疏张量

，其中，

为哈希表，记录着激活位置的坐标，而

为特征矩阵。对于一个稀疏张量，其哈希表和特征矩阵行对行地相互对应。

从这个层面来说，SS-Conv 的目标为用k来卷积输入，

因此可分三个步骤进行实现 SS-Conv：i) 通过旋转可控卷积核的建立，获得卷积核k；ii）通过位置状态的定义，获得输出哈希表

；ii）通过稀疏卷积的操作，获得输出特征矩阵

旋转可控卷积核的建立

满足旋转可控约束条件的关键在于控制特征向量的径向方向，最新的研究发现球形谐波

可以给出解答，基于球形谐波的基核的线性组合来生成旋转可控的卷积核。

首先考虑输入和输出特征都是不可约束的表征，阶数分别为l和k，则在x位置的卷积核

可以表示为基核

的线性组合：

其中，

在上述公式中，

是一组可学习的系数，

是一个连续的高斯径向函数

，

是一个 (2k+1)(2l+1) 大小的基变换矩阵。

再考虑一般情况下，输入和输出特征由不可约的表征堆叠而成，阶数分别为

，则在x位置处整个旋转可控的卷积核k(x)可表示为：

k(x)的大小为

，其中

。

位置状态的定义

SS-Conv 高效性的关键在于位置状态的定义。一般来说，对于一个输出位置x，如果在其感受野中存在任意激活的输入位置，则该位置被激活；否则该位置设为未激活状态，即意味着该位置的特征将被直接设为零向量（基态）。上述定义可表示为：

输出的哈希表即为

。

稀疏卷积的操作

在获得

后，下一个目标为计算

的值。特别地，

先被初始化为零矩阵，继而通过以下算法来更新其中特征向量：

这个过程可分为两小步：第一步是建立规则词典

来记录匹配的激活输入位置y和输出位置x, 第二步是根据R中的匹配关系来更新

。在这个过程中，R的建立非常重要，使得第二步可以在 GPU 上利用矩阵加乘操作高效地实现。

SS-Conv 的性能、速度及显存占用

为了探究 SS-onv 的性能，研究人员在实例层级的 6D 物体姿态估计任务上将其与其他三种 3D 卷积进行比较，包括传统密集连接的 3D 卷积（Dense-Conv）、非 SE(3) 等变的稀疏卷积（SP-Conv）、3D 姿态可控卷积（ST-Conv）。在这些卷积中，SP-Conv 通过考虑数据的稀疏性来提高 Dense-Conv 的速度，在一些 3D 语义任务上非常高效（例如，3D 物体检测）；ST-Conv 建立旋转可控的卷积核，再利用 Dense-Conv 来实现卷积。为了满足不同卷积的多种计算需求，这些实验在相同的实验条件下，基于一个简单轻量的 12 层卷积网络结构（Plain12）上实现。

表 1. 基于不同卷积层的 Plain12 的量化结果。数据库为 LineMOD。

不同卷积在 LineMOD 数据库上的量化结果如表 1 所示，SS-Conv 在准确性和高效性上均占据优势。在准确性方面，SS-Conv 在 ADD(S) 指标上取得了和 ST-Conv 相当的结果，且远远高于 Dense-Conv 及 SP-Conv，充分显示了 SE(3) 等变的特征学习在姿态估计上的重要性。通过逐层地保留特征的相对姿态，SE(3) 等变特性能在特征学习中捕捉更多有效的物体姿态信息。在高效性方面，稀疏姿态可控的卷积神经网络在复杂系统中显得更加快捷和灵活，例如，在 Plain12 中 batch size 为 32 的情况下，SS-Conv 的速度约为 ST-Conv 的 2.7 倍（表 1 中 404FPS v.s. 148FPS）。图 2 展示了更多不同大小的数据 batch 下 FPS 的结果，其中 ST-Conv 在一张最大显存为 12G 的显卡上的极限 batch size 为 48；而 SS-Conv 占据更少显存，即使 batch size 为 512，依然可以运行。值得一提的是，更大的 batch size，SS-Conv 更加高效，例如，在 Plain12 上 batch size 为 512 的情况下，速度可达到 725 FPS。此外，研究人员还在两个更深的网络（Plain24 和 ResNet50）上对 ST-Conv 和 SS-Conv 进行比较，结果与 Plain12 上的保持一致：SS-Conv 比 ST-Conv 享有更快的速度，且占用更少的 GPU 显存。

图 2. 不同 batch size 下 FPS 和 GPU 显存占用曲线。

应用：3D 空间中物体姿态的估计和跟踪

图 3. 基于 SS-Conv 的物体姿态估计的通用框架

如图 3 所示，研究人员基于 SS-Conv 建立了一个两阶段的物体姿态估计框架。在第一阶段中，研究人员首先利用 SS-Conv 建立主干网络，提取多层级的 SE(3) 等变体素特征，后利用 Tenso-to-Point 模块将体素特征转变为观测物体的逐点特征，用于估计初始的物体姿态；在第二阶段中，利用 SS-Conv 的特征空间可控性，研究人员提出了特征操控模块（Feature-Steering module），对等变的体素特征基于预测的初始姿态进行变换，再同样地将其转为逐点特征，用于估计物体残差姿态，从而优化初始预测。第二阶段可被迭代地使用，不断更新物体姿态。

实例层级 6D 物体姿态估计

对于实例层级的任务，研究人员在 LineMOD 数据库上对基于 SS-Conv 的方法同其他已有方法进行对比。不同方法的量化结果如表 2 所示，基于 SS-Conv 的两阶段方法超过目前所有的方法，在 ADD(S) 上取得了新的最高结果 99.2%。可以观察到，得益于 SS-Conv 中特征空间的可控性，基于特征操控模块的第二阶段姿态优化大幅提高了第一阶段的预测结果。