circleNMS源码分析

pla66

已于 2023-06-19 16:44:12 修改

阅读量406

点赞数 1

文章标签： c++

于 2023-05-18 16:18:23 首次发布

本文链接：https://blog.csdn.net/guaaaaaaa/article/details/130744804

版权

文章详细介绍了lidar_centerpoint功能包中使用GPU进行NMS操作的过程。通过circleNMS函数，它利用CUDA在设备上执行计算，将3D框分为多个block，每个block有16个线程，通过核函数circleNMS_Kernel进行距离判断。计算完成后，将结果从设备内存复制到主机内存，生成保留或舍弃掩模。该过程优化了计算效率，适用于大规模3D检测任务。

摘要由CSDN通过智能技术生成

lidar_centerpoint功能包中推理结果的NMS操作是在GPU上计算完成的，下面通过源码来分析一下计算原理。

函数入口

const std::size_t THREADS_PER_BLOCK_NMS = 16;

std::size_t circleNMS(
  thrust::device_vector<Box3D> & boxes3d, 
  const float distance_threshold,
  thrust::device_vector<bool> & keep_mask,
  cudaStream_t stream)  
{
  // num_boxes3d为输入的3d box个数
  const auto num_boxes3d = boxes3d.size();
  // 将所有boxes分为数个block，每个block中有16个thread
  const auto col_blocks = divup(num_boxes3d, THREADS_PER_BLOCK_NMS);
  // 保存每个box与其后方所有boxes的nms结果，每个box有col_blocks个64位mask
  thrust::device_vector<std::uint64_t> mask_d(num_boxes3d * col_blocks);

  // 通过launch函数启动计算
  CHECK_CUDA_ERROR(
    circleNMS_launch(boxes3d, num_boxes3d, col_blocks, distance_threshold, mask_d, stream));

  // memcpy device to host
  thrust::host_vector<std::uint64_t> mask_h(mask_d.size());
  thrust::copy(mask_d.begin(), mask_d.end(), mask_h.begin());
  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));

  // generate keep_mask
  std::vector<std::uint64_t> remv_h(col_blocks);
  thrust::host_vector<bool> keep_mask_h(keep_mask.size());
  std::size_t num_to_keep = 0;
  for (std::size_t i = 0; i < num_boxes3d; i++) {
    auto nblock = i / THREADS_PER_BLOCK_NMS;
    auto inblock = i % THREADS_PER_BLOCK_NMS;

    if (!(remv_h[nblock] & (1ULL << inblock))) {
      keep_mask_h[i] = true;
      num_to_keep++;
      std::uint64_t * p = &mask_h[0] + i * col_blocks;
      for (std::size_t j = nblock; j < col_blocks; j++) {
        remv_h[j] |= p[j];
      }
    } else {
      keep_mask_h[i] = false;
    }
  }

  // memcpy host to device
  keep_mask = keep_mask_h;

  return num_to_keep;
}

其中，函数参数分别为NMS前的3d boxes vector，NMS距离阈值，保存NMS结果的mask vector和cuda stream，前两项为输入，第三项为输出。函数返回NMS后的目标个数。输入的boxes3d vector已按照置信度由高到低的顺序排列。输出的mask为true表示保留，false表示舍弃。

launch函数

cudaError_t circleNMS_launch(
  const thrust::device_vector<Box3D> & boxes3d, 
  const std::size_t num_boxes3d,
  std::size_t col_blocks, 
  const float distance_threshold,
  thrust::device_vector<std::uint64_t> & mask,
   cudaStream_t stream)
{
  // circleNMS是通过计算两个box的中心点距离的平方是否达到阈值的平方来NMS的，此处计算阈值的平方
  const float dist2d_pow_thres = powf(distance_threshold, 2);

  // 开辟col_blocks * col_blocks个block
  dim3 blocks(col_blocks, col_blocks);
  // 每个block开辟16个线程
  dim3 threads(THREADS_PER_BLOCK_NMS);
  // 运行核函数
  circleNMS_Kernel<<<blocks, threads, 0, stream>>>(
    thrust::raw_pointer_cast(boxes3d.data()), num_boxes3d, col_blocks, dist2d_pow_thres,
    thrust::raw_pointer_cast(mask.data()));

  return cudaGetLastError();
}

kernel函数

__global__ void circleNMS_Kernel(
  const Box3D * boxes, 
  const std::size_t num_boxes3d,
   const std::size_t col_blocks,
  const float dist2d_pow_threshold,
   std::uint64_t * mask)
{
  // params: boxes (N,)
  // params: mask (N, divup(N/THREADS_PER_BLOCK_NMS))

  const auto row_start = blockIdx.y;
  const auto col_start = blockIdx.x;

  if (row_start > col_start) return;

  const std::size_t row_size =
    fminf(num_boxes3d - row_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);
  const std::size_t col_size =
    fminf(num_boxes3d - col_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);

  __shared__ Box3D block_boxes[THREADS_PER_BLOCK_NMS];

  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x] = boxes[THREADS_PER_BLOCK_NMS * col_start + threadIdx.x];
  }
  __syncthreads();

  if (threadIdx.x < row_size) {
    const std::size_t cur_box_idx = THREADS_PER_BLOCK_NMS * row_start + threadIdx.x;
    const Box3D * cur_box = boxes + cur_box_idx;

    std::uint64_t t = 0;
    std::size_t start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (std::size_t i = start; i < col_size; i++) {
      if (dist2dPow(cur_box, block_boxes + i) < dist2d_pow_threshold) {
        t |= 1ULL << i;
      }
    }
    mask[cur_box_idx * col_blocks + col_start] = t;
  }
}

为了方便理解核函数的操作，我们假设NMS前一共有24个box，则num_boxes3d=24, col_blocks=2。
我们可以画一个24*24的矩阵，(x, y)表示第x和y个box之间的判断操作。
在这里插入图片描述
图中的蓝色小方块表示对应的box需要进行计算判断，灰色则不需要。
我们将这些小方块按16个一组分为四组：

可以认为这四块代表了四个block中需要进行的计算。每一个block当中的一行小方块代表的计算，由block中的一个thread来完成。可以发现上方的两个block内只有一半thread参与计算，且左上角的block不需要任何计算。
现在我们来看核函数代码，首先，计算当前位于哪一个block：

  const auto row_start = blockIdx.y;
  const auto col_start = blockIdx.x;

  if (row_start > col_start) return;

如果row_start > col_start，即blockIdx.y > blockIdx.x，则此block中的线程不参与计算。对应我们的假设，就是(0, 1)block不参与计算，很合理：
橙色框为block，蓝色框为thread，灰色的block没有计算操作直接返回
然后，需要计算出当前block进行的NMS计算操作都是针对哪些box(row_size)，并且这些box需要针对哪些box(col_size)做计算。将col_size个对应的box放入共享内存中。使用fminf()函数是因为之前的divup操作，有的block中的16个线程并不都需要参与计算，例如这里的(1,1)block。

  // row_size:行数，可以理解为当前block中参与计算的thread数
  const std::size_t row_size =
    fminf(num_boxes3d - row_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);
  // col_size:列数，可以理解为每个thread最多需要进行计算的次数，即前面block中一行小方块的个数
  const std::size_t col_size =
    fminf(num_boxes3d - col_start * THREADS_PER_BLOCK_NMS, THREADS_PER_BLOCK_NMS);
    
  __shared__ Box3D block_boxes[THREADS_PER_BLOCK_NMS];

  // 根据thread索引找到对应的box，存入共享内存
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x] = boxes[THREADS_PER_BLOCK_NMS * col_start + threadIdx.x];
  }

此处每个block的row_size与col_size，与前面图片中block内的小方块个数对应
在这里插入图片描述

在线程同步后，就可以进行计算和判断了：

  __syncthreads();

  if (threadIdx.x < row_size) {
    // cur_box_idx:当前thread对应的box在所有box中的idx
    const std::size_t cur_box_idx = THREADS_PER_BLOCK_NMS * row_start + threadIdx.x;
    // 找到当前box
    const Box3D * cur_box = boxes + cur_box_idx;
    
    // t的每一位代表当前box与block_boxes中相应的box的距离是否小于阈值
    // 此处其实不需要uint64_t，uint16_t即可
    // t可以看做上图中一行16个小方块的0/1状态，注意小方块的idx对应t中由低到高的位数
    std::uint64_t t = 0;
    std::size_t start = 0;
    // 如果是对角线上的block，则block中每一行跳过灰色小方块
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    // 当前box与block_boxes中的box做计算
    for (std::size_t i = start; i < col_size; i++) {
      if (dist2dPow(cur_box, block_boxes + i) < dist2d_pow_threshold) {.
        // 如果两个box距离小于阈值，将t的第i位置1
        t |= 1ULL << i;
      }
    }
    // 将当前线程的t填入mask中相应的位置，每个box有col_blocks个mask
    mask[cur_box_idx * col_blocks + col_start] = t;
  }

核函数看完，接下来回到circleNMS函数：

  CHECK_CUDA_ERROR(
    circleNMS_launch(boxes3d, num_boxes3d, col_blocks, distance_threshold, mask_d, stream));
    
  // memcpy device to host
  thrust::host_vector<std::uint64_t> mask_h(mask_d.size());
  thrust::copy(mask_d.begin(), mask_d.end(), mask_h.begin());
  CHECK_CUDA_ERROR(cudaStreamSynchronize(stream));

  // generate keep_mask
  // remv_h:用于标记box是否应该被舍弃
  std::vector<std::uint64_t> remv_h(col_blocks);
  // keep_mask_h:nms后得到的mask
  thrust::host_vector<bool> keep_mask_h(keep_mask.size());
  std::size_t num_to_keep = 0;
  // 按照boxes3d中的顺序逐个判断box是否应该保留
  for (std::size_t i = 0; i < num_boxes3d; i++) {
    // nblock:当前box位于哪一个block（可以看做blockIdx.y）
    auto nblock = i / THREADS_PER_BLOCK_NMS;
    // inblock:位于block中的哪一个thread
    auto inblock = i % THREADS_PER_BLOCK_NMS;

    // 检查当前box在remv_h中是否被置1，如果为1则舍弃，为0则保留
    if (!(remv_h[nblock] & (1ULL << inblock))) {
      // 当前box将被保留
      keep_mask_h[i] = true;
      num_to_keep++;
      // 找到当前box与其他所有box的NMS结果，即核函数中的t，一个box有col_blocks个t
      // p为指向第一个t的指针
      std::uint64_t * p = &mask_h[0] + i * col_blocks;
      // 从当前block开始，每一个block的t与remv_h按位或后赋值给remv_h，含义为排在当前box后面的box是否应舍弃
      for (std::size_t j = nblock; j < col_blocks; j++) {
        remv_h[j] |= p[j];
      }
    } else {
      // 当前box将被舍弃
      keep_mask_h[i] = false;
    }
  }

  // memcpy host to device
  keep_mask = keep_mask_h;

  return num_to_keep;