手写nms

计算宽高的时候加1是为什么？

本文总结自互联网的多种nms实现，供参考，非博主原创，各原文链接如下，也建议大家动手写一写。

Ref：

浅谈NMS的多种实现

目标窗口检测算法-NMS非极大值抑制

一、faster-rcnn源码阅读：nms的CUDA编程

c++版 nms

nms简介

首先还是要科普一下nms算法的思想：简单来说就是去重框。这里的重框针对的当然是某一类的框。下面实现的时候也是默认拿到某一类所有的框。

算法思路：

For a prediction bounding box B, the model calculates the predicted probability for each category. Assume the largest predicted probability is p, the category corresponding to this probability is the predicted category of B. We also refer to pas the confidence level of prediction bounding box B. On the same image, we sort the prediction bounding boxes with predicted categories other than background by confidence level from high to low, and obtain the list L. Select the prediction bounding box B1 with highest confidence level from L as a baseline and remove all non-benchmark prediction bounding boxes with an IoU with B1 greater than a certain threshold from L. The threshold here is a preset hyper-parameter. At this point,L retains the prediction bounding box with the highest confidence level and removes other prediction bounding boxes similar to it. Next, select the prediction bounding box B2 with the second highest confidence level from L as a baseline, and remove all non-benchmark prediction bounding boxes with an IoU with B2 greater than a certain threshold from L. Repeat this process until all prediction bounding boxes in L have been used as a baseline. At this time, the IoU of any pair of prediction bounding boxes in L is less than the threshold. Finally, output all prediction bounding boxes in the list L.

前面这段话基本说出了算法的具体实现思路：先对每个框的score进行排序，首先选择第一个，也就是score最高的框，它一定是我们要保留的框。然后拿它和剩下的框进行比较，如果IOU大于一定阈值，说明两者重合度高，应该去掉，这样筛选出的框就是和第一个框重合度低的框，第一次迭代结束。第二次从保留的框中选出score第一的框，重复上述过程直到没有框保留了。

Python

版本一

该版本为 Faster RCNN 实现的版本，是网络是最常见的版本：

def nms(dets, thresh):x1 = dets[:, 0] #xminy1 = dets[:, 1] #yminx2 = dets[:, 2] #xmaxy2 = dets[:, 3] #ymaxscores = dets[:, 4] #confidenceareas = (x2 - x1 + 1) * (y2 - y1 + 1) # 每个boundingbox的面积order = scores.argsort()[::-1] # boundingbox的置信度排序keep = [] # 用来保存最后留下来的boundingboxwhile order.size > 0:     i = order[0] # 置信度最高的boundingbox的indexkeep.append(i) # 添加本次置信度最高的boundingbox的index# 当前bbox和剩下bbox之间的交叉区域# 选择大于x1,y1和小于x2,y2的区域xx1 = np.maximum(x1[i], x1[order[1:]]) #交叉区域的左上角的横坐标yy1 = np.maximum(y1[i], y1[order[1:]]) #交叉区域的左上角的纵坐标xx2 = np.minimum(x2[i], x2[order[1:]]) #交叉区域右下角的横坐标yy2 = np.minimum(y2[i], y2[order[1:]]) #交叉区域右下角的纵坐标# 当前bbox和其他剩下bbox之间交叉区域的面积w = np.maximum(0.0, xx2 - xx1 + 1)h = np.maximum(0.0, yy2 - yy1 + 1)inter = w * h# 交叉区域面积 / (bbox + 某区域面积 - 交叉区域面积)ovr = inter / (areas[i] + areas[order[1:]] - inter)#保留交集小于一定阈值的boundingboxinds = np.where(ovr <= thresh)[0]order = order[inds + 1]return keep

可以看到，基本是按照上述思路去写的，按照score进行降序排序，然后每次拿到第一个框，也就是score最大的框，然后计算该框与其他框的IOU，最后留下iou<=thresh的框留作下次循环，这里唯一值得强调的是最后这个索引为什么要+1。这是因为我们要得到的inds是排除了当前用来比较的score最大的框，所以在其原始索引基础上+1 ,从代码中看就是由于order[1:]这样写导致的。

当然，这个大众版本思路很清楚，但感觉不够优雅，可以再精简一点。

版本二

def nms(dets, thresh):                   areas=np.prod(bbox[:,2:]-bbox[:,:2],axis=1)order = scores.argsort()[::-1]keep=[]while order.size>0:i=order[0]keep.append(i)tl=np.maximum(b[:2],bbox[i+1:,:2])br=np.minimum(b[2:],bbox[i+1:,2:])inter=np.prod(br-tl,axis=1)*(br>tl).all(axis=1)ovr=inter/(areas[order[1:]]+areas[i]-inter)inds=np.where(ovr<=thresh)[0]order=order[inds+1]return keep

当然这里的思路还是要排序，只不过iou部分写的更精简了。

好了，铺垫了这么久，说一下不排序怎么写。基本思路是，依次遍历每个框，计算这个框与其他框的iou,找到iou大于一定阈值的其他框，因为这个时候不能保证它一定是score最高的框，所以要进行判断，如果它的score小于其他框，那就把它去掉，因为它肯定不是要保留的框。如果它的score大于其他框，那应该保留它，同时可以去掉所有其他框了。最后保留的框就是结果。

def nms(bbox, scores, thresh):area=np.prod(bbox[:,2:]-bbox[:,:2],axis=1)keep=np.ones(len(bbox),dtype=bool)for i, b in enumerate(bbox):if(keep[i]==False):continuetl=np.maximum(b[:2],bbox[i+1:,:2])br=np.minimum(b[2:],bbox[i+1:,2:])inter=np.prod(br-tl,axis=1)*(br>=tl).all(axis=1)iou=ia/(area[i+1:]+area[i]-inter)r = [ k for k in np.where(iou>thresh)[0]+i+1 if keep[k]==True]if (scores[i]>scores[r]).all():keep[r]=Falseelse:keep[i]=Falsereturn np.where(keep)[0].astype(np.int32)

这是我按照上面思路写的，用keep表示框的去留，为True的框要保留。当然这个思路的效率不见得比前面排序的要好，只是作为其他角度思考。

版本三

补充：之前排序实现是通过每次筛除框来完成的，直到最后没有框剩下了则循环结束。但还可以这么想：同样先将框按score排好序，这次循环条件为遍历所有框，在某次循环，拿到一个框，将其与所有已经保留的框进行比较，如果iou大于阈值，说明它应该被删去，直接进行下一次循环，如果小于将其加入进要保留的框中。当遍历完所有框时结束，拿到保留的框。

嗯，文字还是有点绕，直接看code更清楚：

def _non_maximum_suppression_cpu(bbox, thresh, score):order=np.argsort(score)[::-1]bbox=bbox[order]      bbox_area=np.prod(bbox[:, 2:]-bbox[:, :2],axis=1)keep=np.zeros(bbox.shape[0], dtype=bool)  for i, b in enumerate(bbox):tl=np.maximum(b[:2],bbox[keep, :2])  br=np.minimum(b[2:],bbox[keep, 2:])	   area = np.prod(br-tl, axis=1)*(tl<br).all(axis=1)iou=area/(bbox_area[i]+bbox_area[keep]-area) if (iou>=thresh).any():   continuekeep[i]=Truekeep = np.where(keep)[0]return keep.astype(np.int32)

至此介绍了三种NMS算法的实现思路。即使作为常规的算法题考察也很好，因为这也不牵扯到目标检测的东西。而且还顺便考察了numpy的操作。这也提醒自己有时候要多深入想想，而不是简单copy网上现有的东西。

C++

typedef struct Bbox{int x;int y;int w;int h;float score;
}Bbox;static bool sort_score(Bbox box1,Bbox box2){return box1.score > box2.score ? true : false;
}float iou(Bbox box1,Bbox box2){int x1 = max(box1.x,box2.x);int y1 = max(box1.y,box2.y);int x2 = min(box1.x+box1.w,box2.x+box2.w);int y2 = min(box1.y+box1.h,box2.y+box2.h);// int w = max(0,x2 - x1 + 1);// int h = max(0,y2 - y1 + 1);int w = max(0,x2 - x1 + 1);int h = max(0,y2 - y1 + 1);// float over_area = w*h;float over_area = (x2 - x1) * (y2 - y1);return over_area / (box1.w * box1.h + box2.w * box2.h - over_area);
}vector<Bbox> nms(std::vector<Bbox>&vec_boxs,float threshold){vector<Bbox>results;std::sort(vec_boxs.begin(),vec_boxs.end(),sort_score);while(vec_boxs.size() > 0) {results.push_back(vec_boxs[0]);int index = 1;while(index < vec_boxs.size()){float iou_value = iou(vec_boxs[0],vec_boxs[index]);cout << "iou:" << iou_value << endl;if(iou_value > threshold)vec_boxs.erase(vec_boxs.begin() + index);elseindex++;}vec_boxs.erase(vec_boxs.begin());}return results;
}

测试：

int main(){vector<Bbox> input;Bbox box1 = {1, 2, 2, 2, 0.4};Bbox box2 = {1, 3, 2, 1, 0.5};Bbox box3 = {1, 3, 3, 1, 0.72};Bbox box4 = {1, 1, 3, 3, 0.9};Bbox box5 = {1, 1, 2, 3, 0.45};input.push_back(box1);input.push_back(box2);input.push_back(box3);input.push_back(box4);input.push_back(box5);vector<Bbox> res;res = nms(input, 0.5);for(int i = 0;i < res.size();i++){printf("%d %d %d %d %f",res[i].x,res[i].y,res[i].w,res[i].h,res[i].score);cout << endl;}return 0;
}

CUDA

cpu版验证和理解了算法，下面来看看GPU实现加速，速度大概可以提升50X。在2080it上大概是115ms比3ms。

nms各种实现的benchmark请看：

https://github.com/fmscole/benchmark

主控函数分两部分，第一部分计算mask，还需要的标0，不需要的（重复度大的）标1，第二部分是根据mask，选出留下来的候选框。

先简要说一下CUDA编程模型：

GPU之所以能够加速，是因为并行计算，即每个线程负责计算一个数据，充分利用GPU计算核心超多（几千个）的优势。

（1）每个计算核心相互独立，运行同一段代码，这段代码称为核函数；

（2）每个核心有自己的身份id，线程的身份id是两个三维数组：（blockIdx.x，blockIdx.y，blockIdx.z）-（threadIdx.x，threadIdx.y，threadIdx.z）。

身份id被另两个三维数组grid（gridDim.x,gridDim.y,gridDim.z）和block(blockDim.x,blockDim.y,blockDim.z)确定范围。

总共有 $gridDim.x\times gridDim.y\times gridDim.z$ 个 block

每个block有 $blockDim.x\times blockDim.y\times blockDim.z$ 个 thread。

有了线程的身份id，经过恰当的安排，让身份id（核函数可以获取）：（blockIdx.x，blockIdx.y，blockIdx.z）-（threadIdx.x，threadIdx.y，threadIdx.z）对应到一个数据，就可以实现一个线程计算一个数据，至于如何对应，开发人员得好好安排，可以 说这是CUDA开发的一个核心问题。

gridDim.x、blockIdx.x这些是核函数可以获取的，gridDim.x等于多少，调用核函数的时候就要定一下来。看代码：

  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);

这里的threadsPerBlock=8*8=64，

当boxes_num=12030时，DIVUP(12030, 64)=12030/64+12030%64>0=188

在调用核函数的时候，通过<<<blocks, threads>>>（#这是cu语法，不是标准C语言）把线程数量安排传递进去，核函数里就有

gridDim.x=188,gridDim.y=188,gridDim.z=1；

blockDim.x=64,blockDim.y=1,blockDim.z=1；

0<=blockIdx.x<188,

0<=blockIdx.y<188,

blockIdx.z=0,

0<=threadIdx.x<64，

threadIdx.y=threadIdx.z=0，

这样就启动了2,262,016个（两百多万个线程）来计算，两百多万看起来吓人，对GPU来书毫无负担！每个线程计算不超过64个值，后面再讲。

（3）这里的grid(a,b,c），block(x,y,z)值是多少，由程序设计人员根据问题来定，在调用核函数时就要确定下来，但有一个基本限制block(x,y,z)中的x×y×z<=1024（这个值随GPU版本确定，起码nvidia 1080，2080都是这样）；

（4）block中的线程每32个thread为一束，绝对同步：比如if-else语句，这32个线程中有的满足if条件，有的满足else。满足else的那部分线程不能直接进入，而是要等满足if的那部分线程运行完毕才进入else部分，而满足if的那部分线程现在也不能结束，而是要等else部分线程运行完毕，大家才能同时结束。for语句也是一样。因此GPU计算尽可能不要有分支语句。

不是说不能用if和for，该用还得用，用的时候要知道付出的代价。否则实现了减速都不知道为了啥。

不同的线程束之间不同步，如果同步需要请__syncthreads();

如果设置block(1)，即一个block只安排一个线程呢？事实上GPU还是要启动32个线程，另外31个陪跑。

因此block(x,y,z)中的x×y×z应该为32的倍数。不过32×32=1024了。

（5）要并行计算，前提是数据之间没有相互依赖，有前后依赖的部分只能放在同一个核函数里计算；

先看控制部分代码，这部分做的事情就是：

1. 在GPU上分配内存，把数据传到GPU
2. 调用核函数，计算mask；
3. 把数据传回来，
4. 根据mask把获取保留下来的候选框。

nms_kernel.cu（来源https://github.com/jwyang/faster-rcnn.pytorch）：

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);
//keep_out：返回保留目标框的下标
//num_out：返回保留下来的个数float* boxes_dev = NULL;unsigned long long* mask_dev = NULL;const int col_blocks = DIVUP(boxes_num, threadsPerBlock);
//比如boxes_num=12030时，col_blocks=188，后面都以此为例CUDA_CHECK(cudaMalloc(&boxes_dev,boxes_num * boxes_dim * sizeof(float)));
//在GPU上分配内存CUDA_CHECK(cudaMemcpy(boxes_dev,boxes_host,boxes_num * boxes_dim * sizeof(float),cudaMemcpyHostToDevice));
//把候选框数据复制到GPUCUDA_CHECK(cudaMalloc(&mask_dev,boxes_num * col_blocks * sizeof(unsigned long long)));
//分配mask的内存，用于返回mask的计算结果dim3 blocks(DIVUP(boxes_num, threadsPerBlock),DIVUP(boxes_num, threadsPerBlock));
//(188,188)dim3 threads(threadsPerBlock);
//(64)//调用核函数nms_kernel<<<blocks, threads>>>(boxes_num,nms_overlap_thresh,boxes_dev,mask_dev);std::vector<unsigned long long> mask_host(boxes_num * col_blocks);
//在CPU上分配内存，用于返回mask就算结果CUDA_CHECK(cudaMemcpy(&mask_host[0],mask_dev,sizeof(unsigned long long) * boxes_num * col_blocks,cudaMemcpyDeviceToHost));
//把mask从GPU复制到CPU
//下面这段代码最后再叙述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));
//释放GPU上的内存
}

计算mask的核函数是核心代码，什么是mask呢？从内存里讲是一段连续内存，但应该把它想象成一个矩阵：

比如候选框个数为boxes_num=12030时，由于要计算任意两个候选框之间的IOU是否大于阈值，因此，我们需要建一个12030*12030的矩阵，1表示两个候选框的IOU大于0.7，0表示不大于。这样任意两个候选框之间的关系都可以用这个12030行*12030列（行表示本框，列表示其他框）的矩阵保存。

但是，这里有好几处值得改进的地方：

1）对角线上的值不需要计算，因为意味着自己与自己计算IOU，没意义；

2）上三角与下三角是对称的，只需要用到上三角即可；

3）更重要的是，为了保存0或1的值，真需要建这么大的矩阵吗？事实上每连续的64个0、1刚好构成一个无符号的64位整数（unsigned long long），我们只用一个整数表示即可，这样内存减少64倍！这中间设计到位运算，这不是问题。原本需要一个12030***12030的一个数组来表示这个矩阵，现在只要12030***188的一个unsigned long long数组就可以映射这个12030行，188列的矩阵。

比如：这长度为1230*188数组中的第一组188个整数（相当于矩阵的第一行）记录的是第一个框（按分数排序之后）与其他所有框之间的重叠关系，第一个整数的第一位表示与自己，由于自己与自己不参与计算，所以这个值一定为0.

__global__ void nms_kernel(int n_boxes, float nms_overlap_thresh,float *dev_boxes, unsigned long long *dev_mask) {//n_boxes候选框的个数，比如是12030//nms_overlap_thresh=0.7阈值//dev_boxes候选框//dev_mask存放mask的值const int row_start = blockIdx.y;//把blockIdx.y想象成矩阵的行号，n_boxes=12030时，共有188行，const int col_start = blockIdx.x;//把blockIdx.x想象成矩阵的列号，n_boxes=12030时，共有188列，而每个block有64个线程，188*64=12032// 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);//在block里，row_size和col_size最多取到64，尾部不足64就取余数。先把数据复制一份到共享内存，关于共享内存，后面详述__shared__ float block_boxes[threadsPerBlock * 5];//每个block有64个线程，所以复制64个候选框，每个候选框有4个坐标值和一个分数值，共5个值，//所以每个block分配的共享内存大小为64*5=320if (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) {//每个线程虽然运行的是同一段代码，但看到的表示身份的threadIdx.x是不一样的，row_size值也不一样//具体的说，有188*188个线程看到的threadIdx.x是一样的，因为总共有188*188个block，threadIdx.x的取值范围是0到63const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;//int cur_box_idx这个才是真正的行号，因为grid下面还有block。//CUDA编程就是要把这些表示线程的id对应到具体数据，blockIdx.y和threadIdx.x确定当前候选框const float *cur_box = dev_boxes + cur_box_idx * 5;//取出当前的候选框，别忘了候选框用5个数值表示的，所以要乘5int i = 0;unsigned long long t = 0;//t就是存放64个0、1的整数int start = 0;if (row_start == col_start) {//对角线上的blockstart = threadIdx.x + 1;//自己跟自己就不要计算IOU了}for (i = start; i < col_size; i++) {if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) {//每一个当前框都与其他框计算IOU，（其他框存放在共享内存，是一个复制品）//本线程只负责计算第blockIdx.y×64+threadIdx.x号候选框（当前框）//与blockIdx.x×64～blockIdx.x×64+63这64个候选框（其他框）之间的关系t |= 1ULL << i;}}const int col_blocks = DIVUP(n_boxes, threadsPerBlock);//同上假定下，col_blocks=188dev_mask[cur_box_idx * col_blocks + col_start] = t;//dev_mask总长为12030*188，表示12030行，188列的矩阵//即分为12030段（行），每段188个（列）int64，每个int64标记64个0、1，所以每行的这188个int64可以标记12032个0\1，//就是说第i段（行）的12032个记录的是第i个候选框与其他所有候选框之间的重叠关系。//cur_box_idx表示第几段，每段col_blocks（=188）个数，col_start是本段里第几个数，所以//dev_mask[cur_box_idx * col_blocks + col_start] =t 记录的是第cur_box_idx个候选框//与第col_start×64到（col_start+1）×64-1之间这64个矩形框之间的重叠关系}
}

再总结一下，总共规划了（188,188,1）个block，每个block有（64,1,1）个线程，第（c，r，0）-（t,0,0)号线程负责计算的数据是：r×64+t号候选框与c×64～c×64+63号这64个候选框之间的重叠关系，并存进（r×64+t）×188+c这个整数里。当c，r遍历完188，t遍历完64，这188×188×64个线程就计算完了任意两个候选框之间的关系。

由于cuda里每个线程计算一个数据，相当于cpu里的循环。CPU里的循环是一个一个的算，而cuda里是同时在算。

本质上，这里gird（188,188,1）和block（64,1,1）来代替了三重循环：

for（int i=0；i<188;i++)for（int j=0；j<188;j++)for（int k=0；k<64;k++)。。。。。。

只要把这循环中的i，j，k替换为blockIdx.x，blockIdx.y，threadIdx.x即可。

CUDA编程就是循环替代！——这是我目前的理解。

回过头来看把候选框复制进共享内存这部分，

block_boxes[threadIdx.x * 5 + 0] =dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0];

可以看出这里面只用到了blockIdx.x和threadIdx.x，注意到了没，没用到blockIdx.y，而blockIdx.y的范围是0~187，什么意思，意味着这段代码被重复执行了188次！毕竟，我们启动了188×64×188个线程，而数据只有12030个数据，因此有188个线程执行的是相同的数据！不知道我理解的对不对，望大佬帮我指正。

计算好了mask数组之后，计算保留下来的目标框：

//...................std::vector<unsigned long long> remv(col_blocks);memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks);
//长为188的unsigned long long数组，初始值为0
//这12032个标记位用于记录那些候选框已经被剔除int num_to_keep = 0;
//保留下来的个数for (int i = 0; i < boxes_num; i++) {int nblock = i / threadsPerBlock;int inblock = i % threadsPerBlock;
//把12030个候选框分成188组，每64个一组，第nblock组的第inblock个候选框就是第i个候选框if (!(remv[nblock] & (1ULL << inblock))) {
//还没有被舍弃，这两个条件只要有一个不成立，都意味着没有被剔除keep_out[num_to_keep++] = i;
//把下标记录下来unsigned long long *p = &mask_host[0] + i * col_blocks;
//再回忆一遍，mask_host总长是12030×188
//分为12030段，第i段有188个unsigned long long整数，记录第i个候选框与其他所有候选框之间的信息
//p指向的是第i段的起始位置for (int j = nblock; j < col_blocks; j++) {
//遍历后面的188个整数，总共12032个标记位，多出来的2位不去管它remv[j] |= p[j];
//因为第i个候选框已经保留下来，所以与第i个候选框重复的候选框都标记为去除}}}
//keep_out的前num_to_keep个数就是保留下来的候选框的下标
//.................................