目标检测 | Farthest Point Sampling 及其 CUDA 实现

Farthest Point Sampling 及其 CUDA 实现

目录

• Farthest Point Sampling 及其 CUDA 实现
  • 概述
  • 均匀随机采样
  • Farthest Point Sampling
  • Farthest Point Sampling 的并行版本
    • 参数初始化
    • 子集计算
    • 子集归并

概述

在深度学习中，在mesh模型（网格模型）上直接学习并预测是一个相当复杂的任务，一方面在于没有高效的模型，另一方面在于现实中难以直接获取性质良好的mesh模型。所以我们常会通过采样和扫描的形式将mesh模型或者现实中的物体转化为point cloud（点云），对于一个点云，其定义为一个\(N\times C\)的张量，\(N\)表示有多少个点，\(C\)表示每个点的信息。对于一个三维的点云，一般\(C=3\)，以表示\((x,y,z)\)。如果点云还携带了更多的信息例如反射强度等，则\(C=3+C^\prime\)。而评价这个点云质量的好坏，就是评价这个点云是否尽可能保留了mesh模型中的几何特征。

现在假设我们拥有一个mesh模型，其渲染效果如下图所示：

但是我们的模型不能直接处理这个mesh模型，因此我们通过扫描或者采样的方式得到了一组\(4096\times 3\)的点云：

如果我们凭借肉眼观察，可见这组点云较好地保留了mesh的几何特征，因此这组点云可以称得上是“质量不错”。

但是我们需要把这组\(4096\times 3\)点云输入到一个输入为固定大小\(256\times 3\)的模型中，我们就必须通过下采样将点云采样为\(256\times 3\)。

均匀随机采样

首先我们想到的一种最为朴素的下采样方法就是均匀随机采样（uniform random sampling），这种下采样方法在python中可以被简单实现：

# 所有的点（1 x 4096 x 3）
total_point = ...

# 进行均匀随机采样
num_samples = 256
indices = torch.randint(0, total_point.size(1), (num_samples,))
sampled_points = total_point[0, indices]

# 可视化
plot_pointcloud(sampled_points, title="uniform random sampling")
plt.show()

我们得到采样结果为：

显然这组点云的采样效果并不能称得上“质量不错”，如下图所示，经过采样之后的点并不能很好地保留原始的几何特征，显然均匀随机采样得到的点在分布上却并不均匀。

经过简单分析，这是因为曲率更大的表面，其几何细节更多，但是在均匀随机采样中，高曲率表面的点与低曲率表面的点权重一致，导致了在曲率更大的表面上，采样的失真程度更高。

Farthest Point Sampling

Farthest Point Sampling (FPS)最早被引入深度学习领域应该是Charles R. Qi等人的文章PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space，FPS算法在其中作为PointNet++特征提取的中心点采样层使用。

FPS算法的本质是贪心算法，其基本思想是从点云中选择一个起始点，然后在剩余的点中选择与已选点距离最远的点作为下一个选定点。这个过程迭代进行，直到达到所需数量的采样点为止。通过这种方式，FPS确保了选定的点在整个点云中分布均匀，从而能够较好地表示点云的特征。

FPS的关键步骤包括：

1. 选择起始点： 从点云中随机选择一个点作为起始点。
2. 计算距离： 计算剩余点与已选点之间的距离。
3. 选择最远点： 选择距离已选点最远的点作为下一个选定点。
4. 重复步骤2和步骤3： 重复计算距离和选择最远点的过程，直到达到所需数量的采样点。

如下图所示，在每一次循环迭代的开始，我们有已采样的点集（红色点），尚未被采样的点集（绿色点）。其中我们可以将红色点连接起来，可以看做一个流形（虚线），可以看到距离这个流形表面越远的绿色点，其曲率将会越大，拥有的几何细节也就越多，所以我们在迭代中将距离最远的绿色点加入已采样集合中。通过这种贪心算法的思想，我们确保每次采样都是最高几何细节保留。

我们可以写出每次迭代的数学描述：

\[\begin{align*}
p_1 & = 0 \\
p_k & = \text{argmax}_{p \in \text{P}} \ \min_{p_i \in \text{S}_k} \ d(p_i, p) \\
\text{S}_{k+1} & = \text{S}_k \cup \{ p_k \} \\
\end{align*}
\]

其中\(P\)是点云，为\(N\times C\)的张量，\(S\)是采样得到的点集索引，\(d(x,y)\)是距离函数。

根据上述数学描述，我们可以写出FPS算法的CPU实现：

def farthest_point_sampling(points:torch.Tensor, num_samples:int):
    """
    使用Farthest Point Sampling (FPS)算法从点云中选择一组关键点。

    参数:
    - points (torch.Tensor): 输入点云，每一行是一个点的坐标。
    - num_samples (int): 要选择的采样点的数量。

    返回:
    - selected_points (torch.Tensor): 选择的关键点的坐标。
    """
    num_points = points.size(0)
    selected_indices = points.new_empty(num_samples, dtype=torch.long)
    distances = points.new_ones(num_points) * float('inf')

    # 从输入点云中随机选择一个起始点
    start_index = 0
    selected_indices[0] = start_index
    selected_point = points[start_index]

    # 迭代选择最远点
    for i in range(1, num_samples):
        # 计算每个点与已选点之间的欧几里得距离
        dist_to_selected = torch.norm(points - selected_point, dim=1)

        # 更新距离，选择最远点
        distances = torch.min(distances, dist_to_selected)
        farthest_index = torch.argmax(distances)
        selected_indices[i] = farthest_index
        selected_point = points[farthest_index]

    selected_points = points[selected_indices]
    return selected_points

如下图所示，FPS算法得到结果明显要比均匀随机采样的“质量”更好：

最后，我们再对比一下原始图像（a），采样结果（b），均匀随机采样（c）以及FPS采样（d）：

Farthest Point Sampling 的并行版本

在上一节中，我们给出了FPS算法的数学描述：

\[\begin{align*}
p_1 & = 0 \\
p_k & = \text{argmax}_{p \in \text{P}} \ \min_{p_i \in \text{S}_k} \ d(p_i, p) \\
\text{S}_{k+1} & = \text{S}_k \cup \{ p_k \} \\
\end{align*}
\]

其中\(P\)是点云，为\(N\times C\)的张量，\(S\)是采样得到的点集索引，\(d(x,y)\)是距离函数。

其中

\[p_k = \text{argmax}_{p \in \text{P}} \ \min_{p_i \in \text{S}_k} \ d(p_i, p)
\]

可以被写成并行的版本

\[\begin{align*}
D_k&=\begin{cases}
d_{k1} = \min_{p_i \in \text{S}_k} \ d(p_i, p_1) \\
\cdots\\
d_{kn} = \min_{p_i \in \text{S}_k} \ d(p_i, p_n)\\
\end{cases} & \text{parallel operation}\\
p_k & = \text{argmax}_{p \in \text{P}} D_k\ \\
\end{align*}
\]

注意到\(p_k = \text{argmax}_{p \in \text{P}} D_k\)需要前一步操作的线程同步，所以形成了算法中的一个瓶颈

当然在实际操作中，CUDA核心提供的thread数目有限（单个block的大小有限），为了方便对齐，我们往往会将点云划分为多个子集。每个thread都可以看作一个分治的分支，在内部串行地计算分配给它的子集，最后在线程同步阶段完成所有子集的归并

我们以实际中使用最多的FPS算法CUDA实现为例进行解读，源码可以在这里找到：

https://github.com/erikwijmans/Pointnet2_PyTorch/blob/master/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu

整段代码可以被划分为两个部分，其中第二部分还可以划分为子集计算与子集归并：

• 参数初始化
• 循环采样
  • 子集计算
  • 子集归并

参数初始化

如下述代码所示，这段代码负责为thread开辟内存，计算当前batch id以及确定当前thread的id，并确定起始点（设为0）

  if (m <= 0) return;
  __shared__ float dists[block_size];
  __shared__ int dists_i[block_size];

  int batch_index = blockIdx.x;
  dataset += batch_index * n * 3;
  temp += batch_index * n;
  idxs += batch_index * m;

  int tid = threadIdx.x;
  const int stride = block_size;

  int old = 0;
  if (threadIdx.x == 0) idxs[0] = old;
  __syncthreads();

子集计算

循环采样分为外循环与内循环，外循环是一个\((1,M)\)的循环，其中\(M\)是采样的点数量。

内循负责计算当前迭代的距离，我们为每个thread分配一个子集，其子集的定义为:

\[\text{thread}=\{p|p_{index} \mod S= \text{thread}_\text{id}\}
\]

，其中\(S\)为block的thread数量，\(\text{thread}_\text{id}\)为当前trhread的id。

例如当\(S=512\)时，\(id\)为\(0\)则\(\text{thread}_0=\{0,512,1024,\cdots\}\)，\(id\)为\(1\)则\(\text{thread}_1=\{1,513,1025,\cdots\}，\cdots\)

随后求出每个thread中的最大距离，即：

\[\text{best}_k=\arg\max_{p_\text{index}} \{\min_{p_i \in \text{S}_k} \ d(p_i, p)|p\in\text{thread}_k\}
\]

  for (int j = 1; j < m; j++) {
    int besti = 0;
    float best = -1;
    float x1 = dataset[old * 3 + 0];
    float y1 = dataset[old * 3 + 1];
    float z1 = dataset[old * 3 + 2];
    for (int k = tid; k < n; k += stride) {
      float x2, y2, z2;
      x2 = dataset[k * 3 + 0];
      y2 = dataset[k * 3 + 1];
      z2 = dataset[k * 3 + 2];
      float mag = (x2 * x2) + (y2 * y2) + (z2 * z2);
      if (mag <= 1e-3) continue;

      float d =
          (x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1);

      float d2 = min(d, temp[k]);
      temp[k] = d2;
      besti = d2 > best ? k : besti;
      best = d2 > best ? d2 : best;
    }
    dists[tid] = best;
    dists_i[tid] = besti;
    __syncthreads();
  }

子集归并

完成所有的\(\text{best}_k\)的计算之后，进入子集归并阶段。这里的子集归并是一个硬编码的二路归并，归并规则如下：

\[\text{best}_k=\max(\text{best}_k,\text{best}_{k+n})\quad \text{where}\quad k<n,\ n=256,128,\cdots,1
\]

例如，当thread数量\(S=512\)，第一轮\(n=256\)

算法将归并\(\text{best}_0\)与\(\text{best}_{256}\)，\(\text{best}_1\)与\(\text{best}_{257}\)，\(\text{best}_2\)与\(\text{best}_{258}\)，\(\cdots\)，\(\text{best}_{255}\)与\(\text{best}_{511}\)。

经过第一轮归并后,子集数量减半，只剩下\(256\)个子集。

同理继续归并，直到只剩下最后一个子集\(\text{best}_0\)，将\(\text{best}_0\)的点并入已采样点集中，完成一轮采样。

    if (block_size >= 512) {
      if (tid < 256) {
        __update(dists, dists_i, tid, tid + 256);
      }
      __syncthreads();
    }
    if (block_size >= 256) {
      if (tid < 128) {
        __update(dists, dists_i, tid, tid + 128);
      }
      __syncthreads();
    }
    if (block_size >= 128) {
      if (tid < 64) {
        __update(dists, dists_i, tid, tid + 64);
      }
      __syncthreads();
    }
    if (block_size >= 64) {
      if (tid < 32) {
        __update(dists, dists_i, tid, tid + 32);
      }
      __syncthreads();
    }
    if (block_size >= 32) {
      if (tid < 16) {
        __update(dists, dists_i, tid, tid + 16);
      }
      __syncthreads();
    }
    if (block_size >= 16) {
      if (tid < 8) {
        __update(dists, dists_i, tid, tid + 8);
      }
      __syncthreads();
    }
    if (block_size >= 8) {
      if (tid < 4) {
        __update(dists, dists_i, tid, tid + 4);
      }
      __syncthreads();
    }
    if (block_size >= 4) {
      if (tid < 2) {
        __update(dists, dists_i, tid, tid + 2);
      }
      __syncthreads();
    }
    if (block_size >= 2) {
      if (tid < 1) {
        __update(dists, dists_i, tid, tid + 1);
      }
      __syncthreads();
    }

    old = dists_i[0];
    if (tid == 0) idxs[j] = old;

综上所述，并行版本的FPS算法将时间复杂度从\(O(MN)\)压缩为\(O(M\frac{N}{S})\)，其中\(M\)为采样数目，\(N\)为待采样点云大小，\(S\)为thread数量（单个block的大小）。