cuda——nms_cuda nms

// ------------------------------------------------------------------
// 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
 
 //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);
  float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1);
  return interS / (Sa + Sb - interS);
}
 
//nms kernel
/*
参数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__ 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]; //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;
    // 如果当前边界框所处的block与要比较的边界框所处的block相同，则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) {
        // 1ULL = unsigned long long型的数字1（最高位为第64位）；每一位就代表一个边界框索引，如果大于阈值，则该位置1
        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生效
}
 

//此函数实际上的__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中
  // 存储每个边界框与当前选定的block中的64个边界框的交并比比较情况，用于后续的nms
  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（初始值），即取到box1，box1肯定会保留下来（因为它的置信度最高），即!(remv[nblock] & (1ULL << inblock)) = true一定得成立（故remv的所有元素要初始化为0，原因便在于此）
    // 如果remv[n]中的某一位的值为1，则第n个block中对应的该位所对应的边界框需要被移除，因为该边界框与保留下来的某一边界框的交并比已经超过了所设定的阈值。
    // 当 i = 1时，如果remv[0]的第2位（从1开始）为1，则不进入if，即直接移除不保留；如果为0，则进入if，保留box2的索引，以及更新remv。更新过程就是将box1的dev_mask中的内容（也即当前的remv）与box2的dev_mask中的内容进行按位或，
    // 意思就是如果box3与更新后的remv中的对应为吻合，则我们不需要管是和box1还是box2的交并比超过了阈值，直接将其移除即可。
    if (!(remv[nblock] & (1ULL << inblock))) { 
      keep_out[num_to_keep++] = i; // 如果不大于阈值，则当前边界框应该保留
      // 通过按位或操作来快速形成要移除的边界框索引
      // 判断当前边界框与前面保留下来的边界框之间的交并比是否大于阈值
      
      unsigned long long *p = &mask_host[0] + i * col_blocks; 
      // remv存储当前框与该框所属的block的64个边界框的mask信息，表示是否需要移除
      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));
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167