目标检测中NMS的GPU实现（来自于Faster R-CNN中的nms_kernel.cu文件）

最近要修改Faster R-CNN中实现的GPU版的NMS代码，于是小白的我就看起了CUDA编程，当然也只是浅显地阅读一些教程，快速入门而已，所以具体需要注意的以及一些思想，大家移步此博主的系列教程：

在了解了CUDA编程的核心思想后，我们便可以开始阅读nms_kernel.cu文件了，先直接上源码（部分简单的已经注释），如下：

// ------------------------------------------------------------------
// Faster R-CNN
// Copyright (c) 2015 Microsoft
// Licensed under The MIT License [see fast-rcnn/LICENSE for details]
// Written by Shaoqing Ren
// ------------------------------------------------------------------

#include "gpu_nms.hpp"
#include <vector>
#include <iostream>

//cudaError_t是cuda中的一个类，用于记录cuda错误（所有的cuda函数，几乎都会返回一个cudaError_t）
#define CUDA_CHECK(condition) \
  /* Code block avoids redefinition of cudaError_t error */ \
  do { \
    cudaError_t error = condition; \
    if (error != cudaSuccess) { \
      std::cout << cudaGetErrorString(error) << std::endl; \
    } \
  } while (0)

//DIVUP即实现除法的向上取整
#define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0))

//unsigned long long类型是目前C语言中精度最高的数据类型，为64位精度
//threadsPerBlock即自定义的每个Block所含有的线程数目（每个Block的线程数不宜太多，也不宜太少）
int const threadsPerBlock = sizeof(unsigned long long) * 8; //其实threadsPerBlock = 64


__device__ inline float devIoU(float const * const a, float const * const b) {
  float left = max(a[0], b[0]), right = min(a[2], b[2]);
  float top = max(a[1], b[1]), bottom = min(a[3], b[3]);
  float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f);
  float interS = width * height;
  float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1);
  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1);
  return interS / (Sa + Sb - interS);
}

//nms kernel
/*
参数n_boxes:边界框数目
参数nms_overlap_thresh:交并比阈值
*/
__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
                           const float *dev_boxes, unsigned long long *dev_mask) {
  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

  // if (row_start > col_start) return;

  const int row_size =
        min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
        min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  __shared__ float block_boxes[threadsPerBlock * 5]; //共享内存
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x * 5 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0];
    block_boxes[threadIdx.x * 5 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1];
    block_boxes[threadIdx.x * 5 + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2];
    block_boxes[threadIdx.x * 5 + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3];
    block_boxes[threadIdx.x * 5 + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4];
  }
  __syncthreads(); //同步线程

  if (threadIdx.x < row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const float *cur_box = dev_boxes + cur_box_idx * 5;
    int i = 0;
    unsigned long long t = 0;
    int start = 0;
    if (row_start == col_start) {
      start = threadIdx.x + 1;
    }
    for (i = start; i < col_size; i++) {
      if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) {
        t |= 1ULL << i; //1ULL = unsigned long long型的数字1（最高位为第64位）
      }
    }
    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t;
  }
}

//设置哪个GPU用于nms
void _set_device(int device_id) {
  int current_device;
  CUDA_CHECK(cudaGetDevice(&current_device)); //获取当前GPU序号
  if (current_device == device_id) {
    return;
  }
  // The call to cudaSetDevice must come before any calls to Get, which
  // may perform initialization using the GPU.
  CUDA_CHECK(cudaSetDevice(device_id)); //设置device_id号GPU生效
}

void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
          int boxes_dim, float nms_overlap_thresh, int device_id) {
  _set_device(device_id);

  float* boxes_dev = NULL;
  unsigned long long* mask_dev = NULL;

  const int col_blocks = DIVUP(boxes_num, threadsPerBlock);

  CUDA_CHECK(cudaMalloc(&boxes_dev,
                        boxes_num * boxes_dim * sizeof(float)));
  CUDA_CHECK(cudaMemcpy(boxes_dev,
                        boxes_host,
                        boxes_num * boxes_dim * sizeof(float),
                        cudaMemcpyHostToDevice));

  CUDA_CHECK(cudaMalloc(&mask_dev,
                        boxes_num * col_blocks * sizeof(unsigned long long)));

  dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
              DIVUP(boxes_num, threadsPerBlock));
  dim3 threads(threadsPerBlock);
  nms_kernel<<<blocks, threads>>>(boxes_num,
                                  nms_overlap_thresh,
                                  boxes_dev,
                                  mask_dev);

  std::vector<unsigned long long> mask_host(boxes_num * col_blocks);
  CUDA_CHECK(cudaMemcpy(&mask_host[0],
                        mask_dev,
                        sizeof(unsigned long long) * boxes_num * col_blocks,
                        cudaMemcpyDeviceToHost));

  std::vector<unsigned long long> remv(col_blocks);
  memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);

  int num_to_keep = 0;
  for (int i = 0; i < boxes_num; i++) {
    int nblock = i / threadsPerBlock;
    int inblock = i % threadsPerBlock;

    if (!(remv[nblock] & (1ULL << inblock))) {
      keep_out[num_to_keep++] = i;
      unsigned long long *p = &mask_host[0] + i * col_blocks;
      for (int j = nblock; j < col_blocks; j++) {
        remv[j] |= p[j];
      }
    }
  }
  *num_out = num_to_keep;

  CUDA_CHECK(cudaFree(boxes_dev));
  CUDA_CHECK(cudaFree(mask_dev));
}

1.devIoU()函数

//devIoU计算两个边界框之间的交并比
//__device__是CUDA中的限定词，具体含义如下图
//float const * const a表示a是常量指针常量，即a是一个指针常量（不可修改的指针），指向一个常量
__device__ inline float devIoU(float const * const a, float const * const b) {
  float left = max(a[0], b[0]), right = min(a[2], b[2]);
  float top = max(a[1], b[1]), bottom = min(a[3], b[3]);
  float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f);
  float interS = width * height; //交集
  float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1); //边界框a的面积
  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1); //边界框b的面积
  return interS / (Sa + Sb - interS);
}

限定词	执行（excution）	可调用（callable）	注意事项（notes）
__global__	在设备上执行（GPU）	可由主机（host）调用； 可由计算能力为3的设备调用（callable from the device for devices of compute capability 3）	返回必须为void型（即不能范围任何其他类型）
__device__	在设备上执行（GPU）	只能被设备（device）调用
__host__	在主机上执行（CPU）	只能被主机调用	可以省略

注：

__device__ ：声明一个函数是设备上执行的，仅可以从设备调用；
__global__： 在设备上执行，可以从主机调用；
__host__ ： 声明的函数是在主机上执行的，仅可从主机调用；
__device__和__global__函数不支持递归；
__device__和__global__函数不能声明静态变量在它们内部。

2.nms_kernel()函数

//nms kernel（CUDA编程中的核函数）
/*
参数n_boxes:边界框数目
参数nms_overlap_thresh:交并比阈值
参数dev_boxes:存储边界框信息，每五位组成一个边界框信息，[left.x,left.y,right.x,right.y,class]
参数dev_mask:存储边界框间的交并比是否超过上述阈值的信息，以ULL类型进行表示，与哪个框交并比超过阈值，相应位置1，否则置0（输出参数）
*/
__global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh,
                           const float *dev_boxes, unsigned long long *dev_mask) {
  const int row_start = blockIdx.y; //当前调用的block的y坐标（实际是一个索引）
  const int col_start = blockIdx.x; //当前调用的block的x坐标

  // if (row_start > col_start) return;

  //min()的目的是防止从dev_boxes中读取数据越界（原因是n_boxes不一定被threadsPerBlock整除）
  //实际上只有最后一个block中所需要的线程数目可能小于threadsPerBlock，其余均等于threadsPerBlock
  const int row_size =
        min(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
  const int col_size =
        min(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  //__shared__限定词，即每个block中的所有线程共享内存
  __shared__ float block_boxes[threadsPerBlock * 5]; //数字5即边界框的5个信息
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x * 5 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; //left.x
    block_boxes[threadIdx.x * 5 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; //left.y
    block_boxes[threadIdx.x * 5 + 2] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; //right.x
    block_boxes[threadIdx.x * 5 + 3] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; //right.y
    block_boxes[threadIdx.x * 5 + 4] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; //class
  }
  __syncthreads(); //同步线程(使得当前block中的所有线程均读取到相应边界框信息后再执行后面的代码)

  //以下代码实现某一边界框与其余所有边界框（删去了部分重复）进行交并比的阈值判断
  if (threadIdx.x < row_size) {
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; //当前选中的边界框索引
    const float *cur_box = dev_boxes + cur_box_idx * 5; //当前选中的边界框信息首地址索引
    int i = 0;
    unsigned long long t = 0; //用于记录与当前边界框交并比情况，大于阈值相应位置1
    int start = 0;
    if (row_start == col_start) { //如果当前边界框所处的block与要比较的边界框所处的block相同，则start不从0开始，减少重复计算
      start = threadIdx.x + 1;
    }
    for (i = start; i < col_size; i++) {
      if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) {
        t |= 1ULL << i; //1ULL = unsigned long long型的数字1（最高位为第64位）；每一位就代表一个边界框索引，如果大于阈值，则该位置1
      }
    }
    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);
    dev_mask[cur_box_idx * col_blocks + col_start] = t; //存入当前边界框与当前选定的block中的64个边界框的交并比比较情况，用于后续的nms
  }
}

此函数在理解上可能会有一定困难，以下我以图像的方式稍生动一点来说明该函数在干什么。

A.函数输入的dev_boxes中存储的内容如下图（每一个边界框都有5个信息按顺序存储着）：

B.函数输出的dev_mask中存储的内容如下图（threadsPerBlock即每个block所含有的线程数）：

图中是数字，拿0x11为例说明如下：

0x11 = 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0000 0001 0001（共64位）

其中第5位为1，表示当前矩形框与当前选中的block中的第5号矩形框的交并比大于设定的阈值。

由于0x10处于box1所在的第1位，更具体一步表示就是box1与第一个block中的64个边界框中的第5个（即box5）的交并比大于设定的阈值。

注：上述所有索引都从1开始（算法本身是从0开始）;图中的符号表示向上取整。

C. 具体在干什么

作者的一大巧妙之处是将二维的block中的两维都表示为dev_boxes，即blockIdx.x和blockIdx.y名义上是block的索引，但实际上表示的是将dev_boxes分块后的块索引，如下图：

代码在干的事就是取当前blockIdx.y块中的第threadIdx.x个边界框与当前blockIdx.x块中的所有边界框进行交并比上的判断，由此为后续nms做准备。

但是为了降低部分重复计算，如(box1, box2)和(box2, box1)这成对的重复计算，采用如下代码：

if (row_start == col_start) { //如果当前边界框所处的block与要比较的边界框所处的block相同，则start不从0开始，减少重复计算
    start = threadIdx.x + 1;
}

但细心的你们一定会发现，其实上述代码只避免了相同块中的重复计算，对于不同块之间仍旧存在重复计算，例如(box1, box65)和(box65, box1)，其中box1属于第blockIdx.y = 0块，box65属于blockIdx.y = 1 块。（当然重复计算并不会影响后续的nms，但会消耗时间）

3._nms()函数

//此函数实际上的__host__类型，真正实现nms
/*
参数keep_out:int型指针，用于存储所有保留下来的边界框索引
参数num_out:保留下的边界框数目
参数:boxes_host:输入参数，存储着边界框信息，来自于主机
参数boxes_num:输入的边界框数目
参数boxes_dim:边界框维度（一般为5，即左上角、右下角和类别）
参数nms_overlap_thresh:交并比阈值，用于nms
参数device_id:GPU设备号
*/
void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num,
          int boxes_dim, float nms_overlap_thresh, int device_id) {
  _set_device(device_id); //设置相应设备

  float* boxes_dev = NULL;
  unsigned long long* mask_dev = NULL;

  const int col_blocks = DIVUP(boxes_num, threadsPerBlock); //向上取整，即当前输入分块后的块数目

  CUDA_CHECK(cudaMalloc(&boxes_dev,
                        boxes_num * boxes_dim * sizeof(float))); //开辟显存
  CUDA_CHECK(cudaMemcpy(boxes_dev,
                        boxes_host,
                        boxes_num * boxes_dim * sizeof(float),
                        cudaMemcpyHostToDevice)); //将host输入的数据送入到boxes_dev中

  CUDA_CHECK(cudaMalloc(&mask_dev,
                        boxes_num * col_blocks * sizeof(unsigned long long)));

  dim3 blocks(DIVUP(boxes_num, threadsPerBlock),
              DIVUP(boxes_num, threadsPerBlock)); //所设置的block为二维block，两维的大小相同
  dim3 threads(threadsPerBlock); //每一个block中的线程为一维，均为threadsPerBlock条线程
  nms_kernel<<<blocks, threads>>>(boxes_num,
                                  nms_overlap_thresh,
                                  boxes_dev,
                                  mask_dev); //调用上述定义的核函数获取交并比情况

  std::vector<unsigned long long> mask_host(boxes_num * col_blocks);
  CUDA_CHECK(cudaMemcpy(&mask_host[0],
                        mask_dev,
                        sizeof(unsigned long long) * boxes_num * col_blocks,
                        cudaMemcpyDeviceToHost)); //从device中处理好的数据送回mask_host，进行后续CPU计算

  std::vector<unsigned long long> remv(col_blocks); //存储要移除的边界框索引
  memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); //初始化为0

  //以下正式开始进行nms，思想和CPU版本有所不同，但本质是一样的
  //由于输入此函数的boxes_host是按置信度从高到低排过序，所以第一个边界框肯定会存入keep_out中
  int num_to_keep = 0;
  for (int i = 0; i < boxes_num; i++) {
    int nblock = i / threadsPerBlock; //当前边界框输入哪一个block
    int inblock = i % threadsPerBlock; //当前边界框输入对应block中的第几个

    //当i = 0时，remv[0] = 0（初始值），但由于第一个边界框肯定要存入keep_out中，所以没问题
    if (!(remv[nblock] & (1ULL << inblock))) { //判断当前边界框与前面保留下来的边界框之间的交并比是否大于阈值
      keep_out[num_to_keep++] = i; //如果不大于阈值，则当前边界框应该保留
      unsigned long long *p = &mask_host[0] + i * col_blocks;
      for (int j = nblock; j < col_blocks; j++) {
        remv[j] |= p[j]; //预存入后续所有边界框是否要被移除的信息（相应位为1则移除）
      }
    }
  }
  *num_out = num_to_keep;

  CUDA_CHECK(cudaFree(boxes_dev));
  CUDA_CHECK(cudaFree(mask_dev));
}

此函数的nms部分可能较难理解（越是没几行的代码越是难以理解），我就举个例子引导一下大家的思维：

假如当前的i = 0，即取到box1，根据nms的原理可知，box1肯定会保留下来（因为它的置信度最高），即!(remv[nblock] & (1ULL << inblock)) = true一定得成立（故remv的所有元素要初始化为0，原因便在于此），由此会进入到if中执行里面的代码。

这时关键就来了，作者通过按位或操作来快速形成要移除的边界框索引，即如下代码：

unsigned long long *p = &mask_host[0] + i * col_blocks;
for (int j = nblock; j < col_blocks; j++) {
     remv[j] |= p[j]; //预存入后续所有边界框是否要被移除的信息（相应位为1则移除）
}

所谓的要移除的边界框索引是指：如果remv[n]中的某一位的值为1，则第n个block中对应的该位所对应的边界框需要被移除，因为该边界框与保留下来的某一边界框的交并比已经超过了所设定的阈值。

好了，回到当前的box1，因为所有的边界框都被分配到了相应的块（block）中，所以remv数组的大小为col_blocks，而通过循环按位或后，remv中存储的是box1与其余边界框的交并比比较情况，也即要移除的边界框索引。

当 i = 1时，如果remv[0]的第2位（从1开始）为1，则不进入if，即直接移除不保留；如果为0，则进入if，保留box2的索引，以及更新remv。更新过程就是将box1的dev_mask中的内容（也即当前的remv）与box2的dev_mask中的内容进行按位或，意思就是如果box3与更新后的remv中的对应为吻合，则我们不需要管是和box1还是box2的交并比超过了阈值，直接将其移除即可。

后面的过程依此类推。