关于点云dgcnn中边特征的学习和理解

本篇博文要讨论的是DGCNN中的边特征，作者在语义分割部分提供了边特征的实现过程，代码段如下：

def get_edge_feature(point_cloud, nn_idx, k=20):
  """Construct edge feature for each point
  Args:
    point_cloud: (batch_size, num_points, 1, num_dims)
    nn_idx: (batch_size, num_points, k)
    k: int

  Returns:
    edge features: (batch_size, num_points, k, num_dims)
  """
  og_batch_size = point_cloud.get_shape().as_list()[0]
  point_cloud = tf.squeeze(point_cloud)
  if og_batch_size == 1:
    point_cloud = tf.expand_dims(point_cloud, 0)

  point_cloud_central = point_cloud
  point_cloud_shape = point_cloud.get_shape()
  batch_size = point_cloud_shape[0].value 
  num_points = point_cloud_shape[1].value
  num_dims = point_cloud_shape[2].value

  idx_ = tf.range(batch_size) * num_points
  idx_ = tf.reshape(idx_, [batch_size, 1, 1]) 

  point_cloud_flat = tf.reshape(point_cloud, [-1, num_dims])
  point_cloud_neighbors = tf.gather(point_cloud_flat, nn_idx+idx_)
  point_cloud_central = tf.expand_dims(point_cloud_central, axis=-2) 
  point_cloud_central = tf.tile(point_cloud_central, [1, 1, k, 1]) 

  edge_feature = tf.concat([point_cloud_central, point_cloud_neighbors-point_cloud_central], axis=-1) 
  return edge_feature

这篇文章还是建立在前几篇博客的基础之上，在该专刊下介绍了pairwise_distance计算形状为(B,N,N)的张量，并且根据这个张量利用tf.nn.top_k()进行knn最近邻查找，拿到了形状为(B,N,K)的最近邻索引。
所以 get_edge_feature的参数列表中：
nn_idx是形状为(B,N,K)的最近邻索引
k为20
至于point_cloud，是输入点云的边特征，也是中间层的边特征，形状有两种情况，前一种是(B,N,F,1)，另一种是(B,N,1,64)。

首先，不管输入点云是什么样子，都会通过tf.squeeze()把数量为1的维度剔除。因此点云的形状为(B,N,F)，这里的F是点属性个数或者叫点特征的个数，它可能是一开始点云的各种坐标，法向量等，也可能是网络中间层的特征。
接着，要利用最近邻索引构建一个最近邻点云。
这里先生成一个列向量，从0到B-1，然后再乘以每个批次点的数量N。如下图所示，这个的用途就是标识了每个批次首个点的点号，把这个向量给他reshape为(B,1,1)。

然后，把输入的点云进行拉伸，用reshape拉伸为(B*N,F)的形状，代码中-1代表自动推断。输入进来的点云其实是三个维度，拉伸之后变成两个维度，也就是成为了一个矩阵，其中的结构如下图。

现在到了关键步骤，看着一行代码point_cloud_neighbors = tf.gather(point_cloud_flat, nn_idx+idx_)。

代码中的nn_idx+idx_是形状(B,N,K) + (B,1,1)，因此用到了广播机制，也就是型为(B,1,1)的idx_被广播为B×N×K的张量。此时广播后的idx_很有意思，只有当第0个维度变化时，里面存的数字才变化。举个例子吧，比如[0,i,j]都是0，[1,i,j]都是1×N，[2,i,j]都是2×N，以此类推。这样的目的是什么呢？
因为nn_idx中存储的是元素的索引，所以nn_idx+idx_的结果就可以准确对应索引到上面那个图所示意的近邻点。
现在用tf.gather()这个函数从拉伸过后的点云中，按照nn_idx+idx_索引取切片元素，下面是示意图：

索引的形状是(B,N,K)，根据索引从拉伸的点云矩阵中取出切片元素，它的工作过程是：
第0个Batch的第0行的第0个元素是a，从拉伸过后的点云中取出来第a行
第0个Batch的第0行的第1个元素是b，从拉伸过后的点云中取出来第b行
…
第1个Batch的第0行的第0个元素是a，从拉伸过后的点云中取出来第a行
第1个Batch的第0行的第1个元素是b，从拉伸过后的点云中取出来第b行
…
…
一直取完，因此得到了BNK这么多行的F维度向量。

简而言之（乱写个吧）：
tf.gather()这个方法可以这样理解，它一般有三个参数
tf.gather(参数,索引,轴)
它的目的就是让我们能够仿照“索引”的样子，沿着“轴”在“参数”当中拿元素，构成一个新的东西。
比如本文，tf.gather(point_cloud_flat, nn_idx+idx_),默认的轴是0，索引是(B,N ,K)这样的形状，参数就是已经拉伸成二维矩阵的点云。这个操作告诉我们，即将要仿照(B,N ,K)的形状构造出一个张量。
那好了，现在待构造的张量长得应该和(B,N ,K)差不多，那里面的第一个元素应该是谁呢？挨个来：
第0Batch的第0行，这一行的第0个元素就应该是待构造张量的第一个元素，这个元素就来自于拉伸后矩阵，因为是按照第0轴取，也就是按照行取，从拉伸后矩阵取一行作为第一个元素。但是这个元素它不是一个数字，而是一个向量，没关系，那么我们想要构造的这个张量最终的形状就是(B,N,K,F)了。至于在从拉伸矩阵中取的时候，到底取谁，当然是按照索引来喽。
所以这个tf.gather()操作过后得到的张量形状为(B,N,K,F)，实际上原始的输入点云一个点占据一行，这一行是该点的属性。但是现在变为了，一个点占据K行，对应着包括自己在内的K个近邻点的点属性，管这个张量叫做point_cloud_neighbors，如下图。

现在最难得一部分已经过去了，下面的操作不是很复杂。
原始输入点云的形状为(B,N,F),给它扩充一个维度变成(B,N,1,F),然后用tf.tile()去沿着倒数第二维度复制K份。上面那个图，左边的矩形里面一个点只对应一行属性，然后右面那个一个点对应K行属性，现在的操作就是把左面，每一个点的一行属性给它复制K份，变成和右面一样的形状，管这个张量叫做point_cloud_central。

最后，就是返回point_cloud_central与point_cloud_neighbors-point_cloud_central拼接后的结果。
下面是point_cloud_neighbors-point_cloud_central的示意图

最后是拼接,也是最终的返回值，即边特征,形状为**(B,N,K,2F)**。

验证一下，在tensorflow1.x环境下整俩placeholder，一个是ph_nn_idx作为最近邻索引，一个是ph_point_cloud作为输入点云，输出结果存在res。
ph_nn_idx形状为(B,N,K)
res形状为(B,N,K,F)