LargeKernel3D: Scaling up Kernels in 3D Sparse CNNs
Abstract
近期2D卷积神经网络的进展已经表明大卷积核非常重要。然而，当直接应用大的卷积核在3D卷积神经网络中时，会遇到严重的困难，其中包括2D中成功的模块设计在3D网络中变得惊人的无效，包括流行的深度卷积。为了解决这一重要挑战，我们提出了空间分割卷积及其大卷积核模块。因此，它避免了朴素的3D大卷积核的优化和效率问题。我们的大卷积核3D卷积神经网络，LargeKernel3D，在3D语义分割和目标检测任务中取得了显着的改进。它在ScanNetv2语义分割任务上实现了73.9%的mIoU，而在NDS nuScenes目标检测基准测试上实现了72.8%的NDS，在nuScenes LIDAR排行榜上排名第一。简单的多模态融合可以进一步提高性能，达到74.2%的NDS。此外，LargeKernel3D可以扩展到Waymo 3D目标检测的17×17×17卷积核尺寸。我们首次展示了大卷积核对于3D视觉任务的可行性和必要性。我们的代码和模型可在github.com/dvlab-research/LargeKernel3D上获得。
1. Introduction
3D稀疏卷积神经网络（CNNs）已被广泛应用于3D任务中作为特征提取器，例如语义分割[9,24]和目标检测[16,55,74]。其高效和方便的使用优势确保了它在各种应用中的重要作用，如自动驾驶和机器人技术。然而，最近基于Transformer的方法[45,46,78]从构建有效的感受野的角度挑战了3D稀疏CNN。全局和局部[21,45]自注意机制都能够捕获大范围的上下文信息。2D视觉Transformer（ViTs）也强调了它们在建模长程依赖方面的优势[20,42,51]。相比之下，普通的3D稀疏CNN在这方面存在局限性。这是因为默认的3D稀疏CNN的感受野受到小卷积核大小和稀疏特征的空间分离的限制（由于子流形稀疏卷积[25]的属性）。
有关2D CNN的文献[18,43,62]提出了一系列方法，结合大卷积核来扩大感受野和模型容量。ConvNeXt [43]采用 7 × 7 7 \times 7 7×7深度卷积作为强大的设计，并结合其他训练技术挑战其Swin Transformer对应物[42]。RepLKNet [18]追求极大的 31 × 31 31 \times 31 31×31卷积核大小来提高不同任务的性能。为了确保RepLKNet [18]的有效性，还需要包括深度卷积等其他因素。其他工作[27]也强调了深度卷积的重要性。由于3D任务和2D任务之间的差异，这些方法并不是3D稀疏CNN的好解决方案。
我们首先从两个方面分析了3D大卷积核CNN设计的困难。第一个挑战是效率。很容易理解，3D卷积具有立方体卷积核大小，计算量增加很快。例如，当卷积核从3×3×3变为7×7×7时，模型大小增加了10倍以上。第二个困难存在于优化过程中。数据集可能仅包含数千个场景，无法与2D图像基准[15,40]在规模方面匹配。此外，3D点云或体素是稀疏的，而不是密集的图像。因此，仅优化大卷积核的参数可能不足，导致过拟合。
图1. 具有不同核的稀疏卷积。小核稀疏卷积在局部区域内聚合特征。它是高效的，但由于特征断开和作用范围小而丢失了足够的信息流。大核稀疏卷积能够捕获长程信息，但代价是参数和计算量大。我们提出的空间分区卷积使用大内核大小，在效率上通过在局部邻居之间共享权重。为了简化，我们展示了2D特征。
本文提出了一种新的 3D 大核设计——空间分区卷积。它是一种新的组卷积家族，通过在空间上邻近位置之间共享权重来代替通道级别的组深度卷积[29]。如图1所示，空间分区卷积通过对空间邻居进行分组来重新映射大核(例如7×7)为小核(例如3×3)，同时绝对的大空间大小保持不变。对于效率问题，它只占用少量的模型大小，使参数保持与小核相同。此外，相比于普通的大核对应物，它具有更少的延迟。对于优化挑战，空间维度之间的权重共享使得参数更有机会更新，从而克服了过拟合问题。
本文中，我们提出了空间分区卷积作为3D大卷积核设计的方法。这是一种新的组卷积家族，通过在空间邻近位置共享权重，而不是在通道层级别上进行深度卷积[29]。如图1所示，空间分区卷积通过将大卷积核（例如7×7）通过分组空间邻居重新映射为小卷积核（例如3×3），同时绝对的大空间尺寸保持不变。就效率问题而言，它占用的模型大小很少，使参数保持与小卷积核相同，而且与纯大卷积核的对应方法相比，延迟更低。至于优化挑战，权重在空间维度之间的共享使参数有更多机会更新，从而克服了过拟合问题。为了增加大卷积核[18]的细节捕获能力，我们引入位置嵌入到空间分区卷积中。这对大卷积核大小产生了显著的影响。我们将所提出的块命名为空间分区大卷积核卷积（SW-LK Conv）。如表1所示，我们比较了基线的3D子流形稀疏卷积和我们的效率。随着其卷积核大小的增加，基线的参数和延迟都大幅增加，而我们的方法则更加高效。
SW-LK Conv可以直接替换现有3D卷积网络中的普通卷积层。我们在现有的3D语义分割[9]和目标检测[16、74]网络上建立了大卷积核骨干网络LargeKernel3D。它在小型模型复杂度开销的同时，取得了显著的改进，超过了现有的最先进方法[9、16、74]。大量实验证实了我们在包括ScanNetv2 [13]、nuScenes [4]和Waymo [57]等大规模基准测试中的有效性。在目标检测方面，LargeKernel3D在nuScenes上实现了72.8％的NDS，位居nuScenes激光雷达排行榜首。简单的体素多模态融合方式进一步将其提高到74.2％的NDS，而没有引入太多额外的技巧。更重要的是，它在大规模的Waymo 3D目标检测中可扩展到17×17×17的卷积核尺寸。
2. Related Work
Convolutional Large-kernel Networks 大核设定已经被广泛地用于二维卷积网络中[18, 43, 67]。ConvNeXt [43] 结合了各种训练技巧和大卷积核，以与 Swin Transformer [42] 竞争。RepLKNet [18] 则是一项专注于大核设计的具体工作。它表明，大核尺寸对目标检测和语义分割等任务的益处要大于对图像分类任务的益处。GCN [47] 在语义分割中也表现出了大核设计的改进。还有一些方法通过隐式技术来近似大核尺寸，包括傅里叶域转换[52]和连续函数[53]。其他方法通过扩张卷积[5, 75]、可变形卷积[14, 80]和主动卷积[31]来改变卷积核形状以扩大感受野。这些方案稀疏化了卷积核，很难捕捉到 3D 深度学习中的原始稀疏特征。
3D Feature Extractors 在三维任务[32, 38]中，一个关键的挑战是从稀疏和非均匀的三维几何数据（例如点云）中学习有效的表示。通常有两种类型的三维特征提取器。第一种是直接在点云上学习，使用一系列的PointNet网络[48, 49]。PointNet++[49]中的采样和分组操作也很耗时。后续的一些工作[46, 78]使用网格采样或场景分割来在每次前向运算中节省计算量，并将多次运算的结果组合起来。第二种方法是通过体素化处理点云并应用三维稀疏卷积神经网络。由于其效率优势，这一方法流派已经在各种三维任务中广泛应用，例如三维分割[9-11, 33]和三维物体检测器中的骨干网络[16, 26, 55, 64]。
空间组卷积（SGC）[77]与所提出的空间分区卷积有关。SGC最初是为了3D场景完成任务而设计的。它们都在空间上进行分区操作。不同之处也是根本和清晰的。我们的方法对卷积核进行分区，而SGC将输入特征分成组。就目标而言，SGC更注重通过改变每个空间组中的特征稀疏性来提高效率，而我们的方法则是为了促进3D任务中的大型卷积核。在实验中进行了进一步的分析和比较。
表格1. 普通3D稀疏卷积和我们的SW-LK卷积的效率比较。基准是具有输入和输出通道16的子流形稀疏卷积层。输入数据是一个随机分散的包含80,000个体素的稀疏张量，批处理大小为1。我们在单个2080Ti GPU上测试所有延迟时间，通过平均10次运行结果获得，预热开始时进行额外10次运行。
Vision Transformer 将注意力计算限制在局部区域或窗口中的视觉Transformer与大卷积核模型[18]具有相似的精神。Swin Transformer [42]使用大小为7到12的平移窗口捕获特征。其变种[19, 41]表明，较大的窗口大小对性能有益。Focal Transformer [70]使用自适应补丁大小捕获细粒度局部注意力。SST [21]和Stratified Transformer [34]将基于窗口的自注意力应用于3D检测和分割。
Position Encoding 由于Transformer中的自注意力是置换等价的，因此位置编码被设计用于向网络注入位置信息。它可以分为绝对位置编码和相对位置编码两个流。对于第一组，位置被以精确的方式进行编码。具体来说，编码[63]是由不同频率的正弦函数生成的，然后加到输入中。相比之下，相对位置编码[54]考虑了相对关系。它可以自然地处理更长的序列，具有更好的可扩展性。大多数工作[3、50、65、72]计算令牌之间的相对距离，并使用可学习的转换进行隐式编码。此外，文献[12]表明，填充卷积可以为相对位置编码提供有效的相对位置信息。除了Transformer，最近的工作[71]也表明，位置嵌入对于使用傅里叶变换的CNN是有益的。
鉴于我们的空间分区卷积通过将卷积核分成不同的部分以追求更高的效率，可能会因每个部分的权重共享而导致位置信息模糊。因此，我们使用相对位置编码作为偏置来补充丢失的位置信息。更多细节请参见第4节。
3. Revisiting 3D Sparse CNNs
Preliminary 3D 稀疏卷积由 3D 稀疏卷积和通常的正则和子流形稀疏卷积[25]组成。我们在公式（1）中制定了 3 D 3\mathrm{D} 3D 稀疏卷积。给定一组稀疏输入特征 { x p ∈ P } \left\{\mathrm{x}_{p \in P}\right\} {xp∈P}，其中 c in c{\text {in }} cin 个通道，每个特征 p ∈ P p \in P p∈P 在 3 D 3D 3D 空间中稀疏分布。我们通过卷积处理这些特征，使用核权重 w ∈ R ∣ K ∣ × c in × c out \mathrm{w} \in \mathbb{R}^{|K| \times c_{\text {in }} \times c_{\text {out }}} w∈R∣K∣×cin ×cout 。例如，在 3 D 3D 3D 坐标空间中， w w w 包含大小为 3 和 ∣ K ∣ = 27 |K|=27 ∣K∣=27 的 c in × c out c_{\text {in }} \times c_{\text {out }} cin ×cout 空间核。输出位置 p ˉ \bar{p} pˉ 上的卷积过程表示为：
其中 k k k是从 p ˉ \bar{p} pˉ的3D偏移距离， p ˉ + k \bar{p}+k pˉ+k是围绕中心 p ˉ \bar{p} pˉ的位置。它枚举了所有离散的位置在核空间 K K K中，其中存在稀疏特征 x p ˉ + k \mathrm{x}_{\bar{p}+k} xpˉ+k。
Convolution Types 正常稀疏卷积通过将每个输入位置 p ∈ P p \in P p∈P扩张到卷积核形状 K K K来扩大其输出位置。这大大降低了特征密度并增加了计算负担。因此，每个阶段的第一层仅在下采样时使用步长为2的正常稀疏卷积。相反，子流形稀疏卷积将其输出位置 p ˉ \bar{p} pˉ固定为输入位置 p ∈ P p \in P p∈P，这保留了低层次的计算成本。更多细节请参考原论文[25]。
由于其效率优势，除了下采样层之外，子流形稀疏卷积主宰了大多数3D CNN层。然而，由于局部范围较小，它会错过空间上分离的特征的足够信息流。增加卷积核大小成为缓解这个问题的潜在解决方案。
Obstacles in 3D Large Kernels 3D大卷积核CNN中的首要问题是效率。当我们增加卷积核的大小时，参数数量和计算负担的增长速度比2D CNN要快得多。例如，给定卷积核大小从3增加到7，一个3D卷积层中的参数数量从 27 ⋅ c i n ⋅ c o u t 27\cdot c_{in}\cdot c_{out} 27⋅cin⋅cout增加到 343 ⋅ c i n ⋅ c o u t 343\cdot c_{in}\cdot c_{out} 343⋅cin⋅cout。因此，简单地扩大3D卷积核是不合理的。
图2. 用于3D物体检测的普通3D CNN、普通3D CNN − 2 × -2 \times −2×和我们的LargeKernel3D的有效感受野（ERF）。 普通的3D CNN主干具有不足的ERF大小，几乎无法覆盖附近的汽车。普通3D CNN − 2 × -2 \times −2×是一个更深的版本，其每个阶段的层数增加一倍。更多的层数对ERF的帮助很小，因为稀疏特征之间的断开。我们的LargeKernel3D在物体中心（左）和边缘（右）都获得了大的ERF，捕获了上下文信息。附录中提供了更多插图。建议使用彩色查看并进行放大。
优化难度是第二个障碍。庞大的参数需要足够的数据进行学习。然而，与2D基准相比，3D数据集通常不具备如此大规模。例如，ImageNet [15] 包含数百万张图片，而3D数据集通常只包含不到一千个场景。当在MinkowskiNet-34[9]中将卷积核从 3 × 3 × 3 3 \times 3 \times 3 3×3×3增加到 7 × 7 × 7 7 \times 7 \times 7 7×7×7时，过拟合问题逐渐恶化，性能下降。
2D领域的先前知识对于3D大卷积核CNNs没有帮助。最近的2D方法[18,43]展示了大卷积核CNNs中的一些有用组件，包括深度卷积[29]，层归一化[1]，GELU [28]激活。我们检查了这些受欢迎的组件，并发现它们在3D CNNs中要么无效，要么甚至有害。为了证明这一发现，我们使用MinkowskiNet-34 [9]对ScanNetv2 [13]的3D语义分割数据集进行基准测试和训练。我们将原始激活函数从ReLU更改为GELU，批归一化到层归一化，以及普通卷积到深度操作。直接使用这些组件会导致最终系统在小型和大型内核网络上的性能下降。因此我们需要寻求其他方法来构建3D大卷积核CNNs。
4. 3D Large-kernel Convolutional Network
4.1. Spatial-wise Partition Convolution
标准卷积核可以看作是一个 3 D 3D 3D矩阵，包括输入通道 C i n C_{in} Cin，输出通道 C o u t C_{out} Cout和空间卷积维度 ∣ K ∣ |K| ∣K∣。对于卷积核大小为 k k k，则在 2 D 2D 2D卷积中，空间卷积维度的体积为 ∣ K ∣ = k × k |K|=k \times k ∣K∣=k×k，而在 3 D 3D 3D卷积中为 k × k × k k \times k \times k k×k×k。深度可分离卷积沿着通道维度共享权重，其中组数等于输入通道数。点卷积将卷积核大小固定为 1 1 1，通常作为连接层与深度可分离卷积一起使用，以调整输出通道数。
图3. 空间分区卷积的说明。在普通的大核卷积中，由于输入特征的稀疏性，每个训练步骤中只有很小一部分的权重被涉及和更新。空间分区卷积通过在空间邻居之间共享权重来缓解这个问题。它增加了权重被优化的机会。在推断时，我们将空间分区卷积重新组合成小型卷积，用大范围的特征分配来完成。
与依赖深度可分离卷积提升性能和权衡准确率和计算量之间的关系的2D方法不同，我们经验性地发现深度可分离卷积对于3D任务没有益处，无论是否包含逐点卷积。
我们提出了一种适用于3D大卷积核CNN的空间划分卷积。它在卷积核的空间维度 K K K 上共享权重，而不是在通道维度上共享。它也不同于SGC [77]，后者在输入特征上进行空间分组。特别地，在公式（1）中， w k w_k wk 在局部区域 k ∈ K k \in K k∈K 中共享相同的值。我们因此将原始的 7 × 7 7 \times 7 7×7 大卷积核划分为 3 × 3 3 \times 3 3×3，通过在空间邻居之间共享权重来实现。空间划分卷积是专门为3D任务设计的，具有以下优点，包括效率和性能。
关于效率方面，在 MinkowskiNet-34 [9] 上，当从 3 × 3 × 3 3 \times 3 \times 3 3×3×3 的卷积核 naively 扩大到 7 × 7 × 7 7 \times 7 \times 7 7×7×7 时，模型大小和延迟都会多次大幅增加。相比之下，我们的空间分区卷积避免了参数增加，并且引入的延迟开销也很有限。我们的方法的优势在于稀疏卷积的具体实现。对于一个空间分区卷积，在推理期间直接使用小卷积层，并将其特征分配区域扩大到大卷积核范围。如图3所示，由于权重共享操作，它大大节省了乘法运算，例如从343次降至27次。我们在开源的 spconv 库上实现了这个操作，并在这里省略了细节。表4显示，我们对空间分区卷积的实现比纯大卷积核基线更加高效。空间分区卷积减轻了大卷积核CNNs的优化困难度，使模型大小保持在一个有限的范围内。它避免了需要大量数据进行训练和过拟合问题的出现。
图4.空间方向大核卷积的结构(SW-LK Conv)。它由大核空间分割卷积和可学习的位置嵌入组成。位置嵌入用于弥补大核的细节捕捉能力。
4.2. Kernel-wise Position Encoding
考虑到空间分区卷积是以共享权重的方式设计的，针对空间稀疏性。虽然这种设计是有效的，但仍存在一个问题：一个区域内的体素共享相同的权重，这导致局部细节模糊。随着核大小的增加，这种现象进一步被放大。为了解决这个问题，我们提出了核位置嵌入设计。我们在图4中说明了这个设计。
特别地，我们初始化位置权重 e ∈ R K × c in e \in \mathbb{R}^{K \times c_{\text{in}}} e∈RK×cin，它对应于卷积核。在卷积过程中，我们让输入特征查询等效位置的位置权重并将它们相加。这个过程将式子（1）修改为：
这个简单的设计本质上是通过添加具有相对位置信息的偏置项到输入特征中来解决由权重共享引起的位置不敏感性问题。它几乎没有额外的计算和参数。特别是对于非常大的内核，例如 17 × 17 × 17 17 \times 17 \times 17 17×17×17，它是一个不错的解决方案。有关更多实验结果和详细信息，请参阅第 5 节。
4.3. Large-kernel Architecture
根据上述的设计和观察结果，我们现在描述我们的大卷积核3D CNN的架构。我们将现有的3D骨干网络中的普通3D子流形稀疏卷积[25]层用我们的SW-LK Conv替换，以进行语义分割和目标检测。
3D Semantic Segmentation LargeKernel3D用于3D语义分割的架构类似于MinkowskiNet [9]，包括1个stem层和8个阶段。stem层是 5 × 5 × 5 5 \times 5 \times 5 5×5×5的子流形卷积。前4个阶段用于下采样，后4个阶段用于上采样。对于U-Net结构，前4个特征将连接到上采样阶段。我们用SW-LK Conv替换了MinkowskiNet34 [9]中的原始层，其中空间方向为 7 × 7 × 7 7 \times 7 \times 7 7×7×7的卷积被分为 3 × 3 × 3 3 \times 3 \times 3 3×3×3的块。默认情况下，包括通道和块数在内的其他超参数遵循MinkowskiNet-34 [9]。
3D Object Detection 3D物体检测器中的典型骨干网络 [ 16 , 55 , 74 ] [16,55,74] [16,55,74]由一个stem层和四个阶段组成。 3 × 3 × 3 3 \times 3 \times 3 3×3×3子流形卷积[25]层作为stem。在每个阶段中（除了第一个阶段），都有一个带有步长2的稀疏卷积层进行下采样。我们用SWLK Convs替换其他普通块。对于nuScenes数据集上的模型，我们发现 7 × 7 × 7 7 \times 7 \times 7 7×7×7的卷积核大小已经足够。它们都使用 3 × 3 × 3 3 \times 3 \times 3 3×3×3的空间分组。对于Waymo数据集上的模型，卷积核大小可扩展到 17 × 17 × 17 17 \times 17 \times 17 17×17×17。我们保留骨干网络的最后一个阶段，因为在前几个阶段中已经足够扩大感受野了。我们默认将这些超参数保持不变，与基线检测器 [ 16 , 74 ] [16,74] [16,74]相同。我们在附录中包含了这些详细的数字。
6. Conclusion and Discussion
我们研究了用于3D卷积神经网络的大内核设计，这与2D卷积神经网络中的解决方案有根本区别。我们提出了专门为3D大内核设计的空间划分卷积（SW Conv），它有效地解决了普通3D大内核卷积神经网络中的效率和优化问题。基于这个设计，我们进一步提出了SW-LK Conv和相应的LargeKernel3D，用于3D语义分割和目标检测。我们的3D大内核网络在语义分割和目标检测任务中取得了不错的提升。我们首次展示了3D大内核可以高效有效地实现。我们希望这些发现能够推动3D网络的进一步发展。附录中提供了更多实验和限制性分析。
Limitations LargeKernel3D主要依赖于手工设计的3D语义分割和目标检测基准中的空间核大小。这些大小对于其他数据集或任务可能不是最优的，取决于整体场景大小和数据稀疏性。其他基于搜索的技术[60]可能会有所帮助，我们将在后续进行尝试。
Boarder impacts 所提出的方法可以作为各种3D任务的骨干网络，这可能包括涉及负面社会影响的任务或数据集。
原文链接：https://arxiv.org/pdf/2206.10555.pdf