【三维语义分割模型】PAConv

【版权声明】
本文为博主原创文章，未经博主允许严禁转载，我们会定期进行侵权检索。

参考书籍：《人工智能点云处理及深度学习算法》

本文为专栏《Python三维点云实战宝典》系列文章，专栏介绍地址“【python三维深度学习】python三维点云从基础到深度学习_python3d点云从基础到深度学习-CSDN博客”。配套书籍《人工智能点云处理及深度学习算法》提供更加全面和系统的解析。

1 PAConv卷积结构

PAConv是由香港大学CVMI Lab和牛津大学合作提出的一种基于位置自适应卷积的三维点云处理算法，发表在CVPR 2021《PAConv: Position Adaptive Convolution with Dynamic Kernel Assembling on Point Clouds》，论文地址为“https://arxiv.org/abs/2103.14635”。PAConv的核心思想在于采用动态卷积核来提取点云特征，动态过程由点云的位置来决定。不同点云输入情况下，用于提取点云的卷积核是不同的，即动态可变的。该结构与transformer模型的思路相似，输入数据特征提取所使用的卷积层是可变的，而常规情况下，网路固定后提取特征的卷积核也就相应固定下来。

这种以数据驱动的卷积核构建方式赋予PAConv比2D卷积更多的灵活性，可以更好地处理不规则和无序的点云数据。此外，通过组合权重矩阵而不是从点位置直接预测内核来能够降低模型学习过程的复杂性。其使用方式可通过将常规卷积替换成PAConv即可实现。本节示例程序的模型结构与PointNet++语义分割结构基本一致，重点区别在于SA模块中MLP层的常规卷积被PAConv卷积所取代。模型详细过程可参考前文关于PointNet++介绍的章节。

PAConv卷积会根据点云位置坐标动态调整点云特征提取的MLP层中的卷积核参数。直接效果体现在卷积核参数随输入发生变化，而不是固定不变的，因而能够具有更强的特征提取能力。这一点在可行变卷积（DCN）之中也有类似体现，DCN主要在于动态调整卷积视野范围。共同之处在于卷积是可变的。进一步思考下去，自然语言处理中的transformer将这种位置相关的卷积思路应用得更加完善。因此，该模型中的PAConv也可以替换成与transformer类似的结构。可以明显看到，在性能改善的同时这些方法也会在一定程度上增加数据计算量。

PAConv实现的主要思路是根据点云位置生成调控卷积核参数的权重得分矩阵，然后通过加权操作得到输出结果，主要过程如下所示。

PAConv卷积过程示意图

假设一组点云输入特征维度为Cin，前三维为点云坐标xyz，且分组内点云总数为K。在PointNet++模型中，SA模块分组点云特征提取采用MLP方式提取，MLP为多个卷积全连接层。卷积的输入通道为Cin，输出通道为Cout，那么经过卷积操作后该组点云特征维度为K x Cout。在该模式下，一旦网络参数经训练确认，那么MLP卷积核参数不会随着输入点云的不同而出现变化，即所有点云提取特征的方式是固定的。PAConv卷积的公式推导过程请它的参考论文，本节重点介绍程序实现过程。

PAConv首先根据点云坐标特征xyz_features采用MLP生成卷积核权重得分。MLP的输入xyz_features为点云坐标xyz、点云坐标中心偏差xyz_diff以及点云相对与中心的欧拉距离组成，共7个维度。这说明得分权重与点云自身坐标与中心相对坐标相关联，重点关注点云局部结构特征。MLP配置参数如下所示，卷积核权重得分scores输出维度为npoint x K x M。npoint表示SA采样点数，K表示分组采样点数，M表示卷积核权重参数数量。

Sequential(
  (layer0): ConvModule(
    (conv): Conv2d(7, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (bn): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activate): ReLU(inplace=True)
  )
  (layer1): ConvModule(
    (conv): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activate): ReLU(inplace=True)
  )
  (layer2): ConvModule(
    (conv): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activate): ReLU(inplace=True)
  )
  (layer3): ConvModule(
    (conv): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1))
  )
)
scores = self.mlps(xyz_features)  # (B, M, N, K)，16x1024x32
scores = F.softmax(scores / self.temp_factor, dim=1)#转换为概率权重，16x1024x32
scores = scores.permute(0, 2, 3, 1)  # (B, N, K, M)，1024x32x16

PAConv卷积会从权重参数矩阵weight_bank加权选择出所需卷积核参数。在输入上，点云特征需要减去中心特征并与原特征进行拼接，那么点云输入特征通道数量由Cin转变为2 x Cin。假设点云特征提取卷积核的通道数量为Cout（通道特征作为点云输出特征维度），且每个卷积核权重参数数量为M，那么卷积核权重参数矩阵weight_bank维度为2 x Cin x M x Cout。

点云输入特征与卷积核权重参数矩阵weight_bank得到维度为 npoint x K x M x Cout维度特征。得分矩阵scores与该矩阵相乘进一步得到加权结果，维度为npoints x K x Cout。加权操作是在M维度上进行的，相当于在该维度上将所有参数通过加权合并为卷积核的一个参数。如果我们将K设置为1，即权重矩阵weight_bank的维度设置为2 x Cin x Cout，那么直接可得到分组点云特征，这也是PointNet++中MLP的实现方式。得分矩阵与点云直接特征进行相乘加权求和，将卷积核参数与点云位置坐标相关联，在一定程度上使得点云中不同目标获得了不同的特征提取方式。

根据上述分析可知，SA模块MLP中的常规卷积被替换成PAConv卷积，用PAConv(Cin, K, M , Cout)来表示，那么得分矩阵scores输出维度为npoint x K x M，卷积核权重参数矩阵weight_bank维度为2 x Cin x M x Cout。输出分组点云特征维度为npoints x K x Cout，即Cout x npoints x K。所起到的最终效果与PointNet++的MLP卷积层基本一致。

2 主干网络

PAConv模型的主干网络采用PointNet2SASSG结构，输入点云数据维度为4096x9。9个属性维度介绍请参考PointNet++章节。网络中包括4个SA模块，每个SA模块中均包含一个由三组PAConv卷积组成的MLP层。

2.1 SA1

采样方式DFPS，采样点数1024，分组采样点数32（KNN）。采样后点云输入坐标point_xyz的维度为3x1024x32，点云输入特征features维度为9x4096。PAConv通过point_xyz和MLP得到卷积核权重分数scores，维度为1024x32x16（K=32，M=16）。

经过采样后点云特征维度为9x1024x32，减去采样中心点特征后得到新特征并与原特征进行拼接，从而得到18x1024x32维特征，即以1024个采样点为中心进行分组，每个组内点云输入特征features由18个维度组成，即2 x Cin。

点云输入特征直接与权重矩阵weight_bank相乘即可得到新的点云特征，即直接使用MLP操作。此时点云输入特征features维度为1024x32x18，权重矩阵weight_bank的维度为18x512（16x32），相乘后新的点云特征new_features维度为1024x32x16x32（npoint，K，M，Cout）。

将权重得分Scores与新的点云特征new_features相乘，PAConv实现了各个卷积核的加权操作，并进一步得到维度为32x1024x32维度的特征。加权操作是在M维度上进行的，相当于在该维度上将所有参数通过加权合并为卷积核的一个参数。如果我们在上一步中将K设置为1，即权重矩阵weight_bank的维度设置为18x32，那么直接可得到分组点云特征，这也是PointNet++中MLP的实现方式。该32x1024x32维度的新特征作为SA模块的分组点云特征。

32x1024x32维度的分组点云特征经过最大池化后得到分组特征，维度为32x1024。以该特征再次进行PAConv操作，此时Cin、K、M、Cout分别为32、32、16、32，输出点云特征维度为32x1024x32，经过最大池化后同样得到32x1024。该特征再次进行一次类似PAConv操作和池化操作，且Cout为64，得到SA1的分组特征，维度为64x1024。

根据上述分析可知，SA模块MLP中的常规卷积被替换成PAConv卷积，并通过连续三次卷积操作完成更深层特征提取。这里用PAConv(Cin, K, M , Cout)来表示，三个卷积参数依次为PAConv(9, 32, 16 , 32)、PAConv(32, 32, 16 , 32)、PAConv(32, 32, 16 , 64)。分组点云特征维度为64x1024x32，经过最大池化得到64x1024维度分组特征。

关键函数解析如下所示。

features, points_xyz = inputs#9x1024x32，Cin=9，npoint=1024，K=32
B, _, npoint, K = features.size()#1024x32
center_features = features[..., :1].repeat(1, 1, 1, K)#中心特征，即采样点特征
features_diff = features - center_features#特征差值
features = torch.cat((features_diff, features), dim=1)#18x1024x32，2 x Cin x npoint x K
xyz_features = self._prepare_scorenet_input(points_xyz)#center_xyz、xyz_diff、euclidian_dist，7x1024x32，7 x npoint x K
scores = self.scorenet(xyz_features)  # [B, npoint, K, M]#1024x32x16
new_features = torch.matmul(features.permute(0, 2, 3, 1),#1024x32x18 18x512
self.weight_bank).view(B, npoint, K, self.num_kernels,#1024x32x512 1024x32x16x32
                                   -1)  # [B, npoint, K, M, Cout]，1024x32x16xCout

new_features = assign_score(scores, new_features)#1024x32xCout
new_features = new_features.permute(0, 3, 1, 2).contiguous()#
new_features = new_features.permute(0, 3, 1, 2).contiguous()#32x1024x32，Cout x npoint x K

2.2 SA2

采样方式DFPS，采样点数256，分组采样点数32（KNN）。

三个卷积参数依次为PAConv(67, 32, 16 , 64)、PAConv(64, 32, 16 , 64)、PAConv(64, 32, 16 , 128)。分组点云特征维度为128x256x32，经过最大池化得到128x256维度分组特征。

2.3 SA3

采样方式DFPS，采样点数16，分组采样点数32（KNN）。

三个卷积参数依次为PAConv(131, 32, 16 , 128)、PAConv(128, 32, 16 , 128)、PAConv(128, 32, 16 , 256)。分组点云特征维度为256x64x32，经过最大池化得到256x64维度分组特征。

2.4 SA4

采样方式DFPS，采样点数64，分组采样点数32（KNN）。

三个卷积参数依次为PAConv(259, 32, 16 , 256)、PAConv(256, 32, 16 , 256)、PAConv(256, 32, 16 , 512)。分组点云特征维度为5121x16x32，经过最大池化得到512x16维度分组特征。

2.5 SA模块完整输出

SA模块全部输出如下所示。

采样点坐标sa_xyz：[xyz_0, xyz_1, xyz_2, xyz_3, xyz_4]，4096x3, 1024x3, 256x3, 64x3, 16x3]

分组特征sa_features：[features_0, features_1, features_2, features_3, features_4]，6x4096, 64x1024, 128x256, 256x64, 512x16]

采样点索引sa_indices：[indices_0, indices_1, indices_2, indices_3, indices_4]，[4096, 1024, 256, 64, 16]

3 特征上采样与结果预测

特征上采样过程与PointNet++结构一致，输出特征fp_feature维度为128x4096。这相当于为每个输入点提取到128维特征，并用该特征来预测语义类别。

结果预测

fp_feature（128x4096）经过卷积Conv1d(128, 128, kernel_size=(1,), stride=(1,), bias=False)和Conv1d(128, 13, kernel_size=(1,), stride=(1,))得到13x4096维预测结果。每个点分别对13个类别进行预测。

参数配置如下所示。

PAConvHead(
  (loss_decode): CrossEntropyLoss(avg_non_ignore=False)
  (conv_seg): Conv1d(128, 13, kernel_size=(1,), stride=(1,))
  (dropout): Dropout(p=0.5, inplace=False)
  (FP_modules): ModuleList(
    (0): PointFPModule(
      (mlps): Sequential(
        (layer0): ConvModule(
          (conv): Conv2d(768, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (activate): ReLU(inplace=True)
        )
        (layer1): ConvModule(
          (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)

4 损失函数

PAConv模型损失包含语义分割损失和PAConv参数正则化损失两部分。

语义预测损失函数为CrossEntropyLoss，其输入为13x4096维预测结果和4096个语义标签pts_semantic_mask。

由于权重矩阵是随机初始化的并且可能会收敛为彼此非常相似的结果，因此无法保证权重矩阵的多样性。为了避免这种情况，作者设计了一个权重正则化损失函数来惩罚不同权重矩阵之间的相关性，以使不同的权重矩阵将更分散和独立，进一步确保了所生成卷积核的多样性。

4个SA结构中分别包含了3个PAConv卷积，因而PAConv参数正则化损失由12部分求和得到，并且乘以10倍权重得到最终损失。PAConv参数正则化损失的定义如下所示。

5 顶层结构

PAConv模型顶层结构主要包括以下三部分：

主干网络特征提取，采用PointNet2SASSG结构提取五组不同尺度特征，见第（2）部分。
语义分割结果预测，包括特征上采样、语义预测和语义分割损失函数计算，见第（3）和（4）部分。
PAConv参数正则化损失，见第（4）部分。

def forward_train(self, points, img_metas, pts_semantic_mask):
    points_cat = torch.stack(points)#输入点云坐标及属性 4096x9
    pts_semantic_mask_cat = torch.stack(pts_semantic_mask) #输入点云语义类别标签，即每个点所属类别
    # extract features using backbone
    x = self.extract_feat(points_cat)
    losses = dict()
    loss_decode = self._decode_head_forward_train(x, img_metas, pts_semantic_mask_cat)
    losses.update(loss_decode)
    if self.with_auxiliary_head:
        loss_aux = self._auxiliary_head_forward_train(
            x, img_metas, pts_semantic_mask_cat)
        losses.update(loss_aux)
    if self.with_regularization_loss:
        loss_regularize = self._loss_regularization_forward_train()
        losses.update(loss_regularize)
return losses

6 模型训练

PAConv模型官方程序地址为“https://github.com/CVMI-Lab/PAConv”，而本节基于mmdetection3d框架中的实现程序进行介绍，其输入数据集为S3DIS。模型训练命令为“python tools/train.py configs/paconv/paconv_cuda_ssg_8x8_cosine_200e_s3dis_seg-3d-13class.py”。运行训练命令可得到如图所示训练结果。