特征聚合模块被并行应用到每个送进网络的点,它包含三个神经单元:局部空间编码、注意力池化、扩张残差块。
1、局部空间编码
给定一个带有每个点特征(特征比如原始RGB、中间学习到的特征)的点云P,这个局部空间编码显式地嵌入全部邻近点的x-y-z坐标,这样的话相应的点特征总是知道它们的相对空间位置。这允许LocSE单元显式地观察局部几何模式,从而最终有利于整个网络有效地学习复杂的局部结构。特别地,本单元包括以下步骤:
找邻近点:对于第i个点,为了提高效率,首先采用简单的KNN算法对其邻近点进行收集。KNN是基于逐点欧氏距离的。
相对点位置编码:对于中心点pi的每一个最近的K点{pi1··pik··piK},我们显式地将相对点位置编码如下:
其中:pi和pik是点的x-y-z位置,⊕是连接操作,|| ||计算邻近点和中心点之间的欧氏距离,似乎rik是从多余的点位置编码而来的。有趣的是,这有助于网络学习局部特征,并在实践中取得了良好的性能。
点特征增强:对于每个相邻点pik,将编码后的相对点位置rik与其对应的点特征fik进行连接,得到增强特征向量fik尖。
最后,相对位置编码单元的输出是一个邻近点特征的新集合Fi尖,它显式地编码了中心点pi的局部几何结构。我们注意到最近的工作[36RS-CNN]也使用点位置来改进语义分割。然而,在[36]中,位置用于学习点得分,而我们的LocSE显式编码相对位置以增强相邻点特征。
# 相对位置编码,输入xyz坐标和邻近点序号
def relative_pos_encoding(self, xyz, neigh_idx):
neighbor_xyz = self.gather_neighbour(xyz, neigh_idx) #邻近点的特征信息
xyz_tile = tf.tile(tf.expand_dims(xyz, axis=2), [1, 1, tf.shape(neigh_idx)[-1], 1])
# 和neighbor_xyz转成一样的数据格式
relative_xyz = xyz_tile - neighbor_xyz
relative_dis = tf.sqrt(tf.reduce_sum(tf.square(relative_xyz), axis=-1, keepdims=True)) # 邻近点与中心点距离
relative_feature = tf.concat([relative_dis, relative_xyz, xyz_tile, neighbor_xyz], axis=-1)
return relative_feature
# 收集相邻点的坐标和特征
@staticmethod
def gather_neighbour(pc, neighbor_idx):
batch_size = tf.shape(pc)[0]
num_points = tf.shape(pc)[1]
d = pc.get_shape()[2].value
index_input = tf.reshape(neighbor_idx, shape=[batch_size, -1])
features = tf.batch_gather(pc, index_input)
features = tf.reshape(features, [batch_size, num_points, tf.shape(neighbor_idx)[-1], d])
return features
2、注意力池化
该神经单元用于聚合相邻点特征集。现有的工作[44pointnet++,33pointcnn]通常使用最大/均值池很难整合相邻的特征,导致大部分信息丢失。相比之下,我们求助于强大的注意机制来自动学习重要的局部特征。特别是,受[65]的启发,我们的注意力池化单元包括以下步骤:
计算注意力分数。给定一个局部特征集合Fi尖,设计一个共享函数g()来学习每个特征独特的注意力分数。基本上,函数g()由一个共享的MLP和softmax组成。它的正式定义如下:
这里W是共享MLP的可学习权值。
加权求和。学习到的注意力分数可以看作是一个自动选择重要特征的软掩膜。形式上,这些特征被加权求和如下:
总之,给定输入点云P,对于第i个点pi,我们的局部位置编码和注意力池化单元学习聚合其K个最近点的几何模式和特征,并最终生成一个信息特征向量fi~ 。
@staticmethod
def att_pooling(feature_set, d_out, name, is_training):
batch_size = tf.shape(feature_set)[0]
num_points = tf.shape(feature_set)[1]
num_neigh = tf.shape(feature_set)[2] # 邻近点个数
d = feature_set.get_shape()[3].value
f_reshaped = tf.reshape(feature_set, shape=[-1, num_neigh, d])
att_activation = tf.layers.dense(f_reshaped, d, activation=None, use_bias=False, name=name + 'fc')
att_scores = tf.nn.softmax(att_activation, axis=1)
f_agg = f_reshaped * att_scores
f_agg = tf.reduce_sum(f_agg, axis=1) # 公式(3)
f_agg = tf.reshape(f_agg, [batch_size, num_points, 1, d])
f_agg = helper_tf_util.conv2d(f_agg, d_out, [1, 1], name + 'mlp', [1, 1], 'VALID', True, is_training) # 共享MLP
return f_agg
(3)扩张残差块
由于点云将被大量下采样,因此我们希望显著增加每个点的感受野,这样即使删除一些点,输入点云的几何细节也更有可能被保留。如图3所示,受成功的ResNet[19]和有效的扩展网络[13]的启发,我们将多个LocSE和注意力池化单元与一个跳过连接堆叠起来,作为扩展剩余块。
为了进一步说明扩展残块的能力,图4显示红色3D点在第一个LocSE/注意力池操作后观察K个相邻点,然后能够从K2相邻点接收信息,即在第二次操作后它的两跳邻域。这是一种通过特征传播扩大感受域和扩大有效邻域的廉价方法。从理论上来说,我们堆叠的单位越多,这个方块就越强大,因为它的触及范围会越来越大。然而,更多的单元必然会牺牲整体的计算效率。此外,整个网络很可能过度贴合。在我们的RandLA-Net中,我们简单地将两组LocSE和注意力池化作为标准残差块进行堆叠,在效率和有效性之间取得了令人满意的平衡。
总体而言,我们的局部特征聚合模块旨在通过显式考虑邻近几何图形和显著增加感受野,有效地保存复杂的局部结构。此外,该模块仅由前馈MLP组成,因此计算效率高。
def dilated_res_block(self, feature, xyz, neigh_idx, d_out, name, is_training):
f_pc = helper_tf_util.conv2d(feature, d_out // 2, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training) # 经过第一个共享MLP
f_pc = self.building_block(xyz, f_pc, neigh_idx, d_out, name + 'LFA', is_training)
f_pc = helper_tf_util.conv2d(f_pc, d_out * 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training,
activation_fn=None)
shortcut = helper_tf_util.conv2d(feature, d_out * 2, [1, 1], name + 'shortcut', [1, 1], 'VALID',
activation_fn=None, bn=True, is_training=is_training)
return tf.nn.leaky_relu(f_pc + shortcut)
def building_block(self, xyz, feature, neigh_idx, d_out, name, is_training):
d_in = feature.get_shape()[-1].value
f_xyz = self.relative_pos_encoding(xyz, neigh_idx)
f_xyz = helper_tf_util.conv2d(f_xyz, d_in, [1, 1], name + 'mlp1', [1, 1], 'VALID', True, is_training)
f_neighbours = self.gather_neighbour(tf.squeeze(feature, axis=2), neigh_idx)
f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
f_pc_agg = self.att_pooling(f_concat, d_out // 2, name + 'att_pooling_1', is_training)
f_xyz = helper_tf_util.conv2d(f_xyz, d_out // 2, [1, 1], name + 'mlp2', [1, 1], 'VALID', True, is_training)
f_neighbours = self.gather_neighbour(tf.squeeze(f_pc_agg, axis=2), neigh_idx)
f_concat = tf.concat([f_neighbours, f_xyz], axis=-1)
f_pc_agg = self.att_pooling(f_concat, d_out, name + 'att_pooling_2', is_training)
return f_pc_agg