CUDA-NMS

CUDA-NMS

这是用cuda写nms的一个例子。

首先用python回忆一下nms的大致流程,

import numpy as np

def iou(x, ys):
    # x single
    # ys mul
    lx = np.maximum(x[0], ys[:,0])
    ly = np.maximum(x[1], ys[:,1])
    rx = np.minimum(x[2], ys[:,2])
    ry = np.minimum(x[3], ys[:,3])

    I = np.maximum(0, rx-lx+1)*np.maximum(0, ry-ly+1)
    U = (x[2]-x[0]+1)*(x[3]-x[1]+1) +(ys[:,2]-ys[:,0]+1)*(ys[:,3]-ys[:,1]+1) - I
    return I*1.0/U

def nms(dets, thresh=0.5):
    # preds: score, bbox[4] (n,5)
    x = dets[:,0]
    y = dets[:,1]
    x1 = dets[:,2]
    y1 = dets[:,3]
    # areas
    areas = 1.0*(x1-x+1)*(y1-y+1)   #n

    order = dets[:,4].argsort()[::-1]
    keep = []
    while order.size >0:
        i = order[0]
        keep.append(i)
        # cal IOU with others
        ious = iou(dets[i,:], dets[order[1:],:])
        passed = np.where(ious<=thresh)[0]
        order = order[passed+1]
    return dets[keep,:]

大致过程是上面是算iou的,下面的是nms,传入的是检测框,和阈值,整个过程就是先选出分值最高的那个, 然后从里面去掉和其相交太多的。不断地重复这个过程。

gpu的

这个是看的网上的代码,对着理解了一下,感觉写的比较高深。同时也对pytorch如何调用c++和cuda的代码有了一定的认识。实践出真知,一定要多动手。

现在的目录结果是这样的,

.
├── build.py
├── _ext
│   ├── __init__.py
│   ├── nms
│   │   ├── __init__.py
│   │   ├── _nms.so
│   │   └── __pycache__
│   │       └── __init__.cpython-35.pyc
│   └── __pycache__
│       └── __init__.cpython-35.pyc
├── __init__.py
├── nms_wrapper.py
├── pth_nms.py
├── __pycache__
│   ├── __init__.cpython-35.pyc
│   ├── nms_wrapper.cpython-35.pyc
│   └── pth_nms.cpython-35.pyc
└── src
    ├── cuda
    │   ├── nms_kernel.cu
    │   ├── nms_kernel.cu.o
    │   └── nms_kernel.h
    ├── nms.c
    ├── nms_cuda.c
    ├── nms_cuda.h
    └── nms.h

其中_ext这个目录是由python build来生成的。当外界调用nms的时候,用的是from nms.nms_wrapper import nms, 其中 nms_wrapper.py这个代码我理解就是一个装饰,因为里面的内容很容易。

    from __future__ import absolute_import
  8 from __future__ import division
  9 from __future__ import print_function
 10 
 11 from nms.pth_nms import pth_nms
 12 
 13 
 14 def nms(dets, thresh):
 15   """Dispatch to either CPU or GPU NMS implementations.
 16   Accept dets as tensor"""
 17   return pth_nms(dets, thresh)


所以除了cuda代码之外,最重要的就是pth_nms.py了,这个里面给了cpugpu两种选择。

import torch
from ._ext import nms
import numpy as np

def pth_nms(dets, thresh):
  """
  dets has to be a tensor
  """
  if not dets.is_cuda:
    x1 = dets[:, 1]
    y1 = dets[:, 0]
    x2 = dets[:, 3]
    y2 = dets[:, 2]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.sort(0, descending=True)[1]
    # order = torch.from_numpy(np.ascontiguousarray(scores.numpy().argsort()[::-1])).long()
    keep = torch.LongTensor(dets.size(0))
    num_out = torch.LongTensor(1)
    nms.cpu_nms(keep, num_out, dets, order, areas, thresh)

    return keep[:num_out[0]]
  else:
    x1 = dets[:, 1]
    y1 = dets[:, 0]
    x2 = dets[:, 3]
    y2 = dets[:, 2]
    scores = dets[:, 4]

    dets_temp = torch.FloatTensor(dets.size()).cuda()
    dets_temp[:, 0] = dets[:, 1]
    dets_temp[:, 1] = dets[:, 0]
    dets_temp[:, 2] = dets[:, 3]
    dets_temp[:, 3] = dets[:, 2]
    dets_temp[:, 4] = dets[:, 4]
    

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.sort(0, descending=True)[1]   #得分从高到低的id。这里好像都一样了。
    dets = dets[order].contiguous()
    keep = torch.LongTensor(dets.size(0))   # 6000
    num_out = torch.LongTensor(1)  # construct a 1 Long tensor, 1 is shape,这两个都是随机的。因为在cuda核里面会改变
    nms.gpu_nms(keep, num_out, dets_temp, thresh)   # this call cuda kernel function
    return order[keep[:num_out[0]].cuda()].contiguous()   # in s


可见无论是用cpu还是用gpu都在_ext中写好了。其实在_ext中的是可执行文件。 原码是在src下面,里面有分别对应cpu的和对应gpu的。

先看看cpu的,记得当时在某云面试的时候就让我用c++来写nms,但是当时我只会用python写。

#include <TH/TH.h>
#include <math.h>

int cpu_nms(THLongTensor * keep_out, THLongTensor * num_out, THFloatTensor * boxes, THLongTensor * order, THFloatTensor * areas, float nms_overlap_thresh) {
    // boxes has to be sorted
    THArgCheck(THLongTensor_isContiguous(keep_out), 0, "keep_out must be contiguous");
    THArgCheck(THLongTensor_isContiguous(boxes), 2, "boxes must be contiguous");
    THArgCheck(THLongTensor_isContiguous(order), 3, "order must be contiguous");
    THArgCheck(THLongTensor_isContiguous(areas), 4, "areas must be contiguous");
    // Number of ROIs
    long boxes_num = THFloatTensor_size(boxes, 0); 
    long boxes_dim = THFloatTensor_size(boxes, 1); 

    long * keep_out_flat = THLongTensor_data(keep_out);
    float * boxes_flat = THFloatTensor_data(boxes);
    long * order_flat = THLongTensor_data(order);
    float * areas_flat = THFloatTensor_data(areas);

    THByteTensor* suppressed = THByteTensor_newWithSize1d(boxes_num);
    //这个是记录每个bbox的状态的。即有没有被前面的box给过滤掉。
    THByteTensor_fill(suppressed, 0); 
    unsigned char * suppressed_flat =  THByteTensor_data(suppressed);

    // nominal indices
    int i, j;
    // sorted indices
    int _i, _j; 
    // temp variables for box i's (the box currently under consideration)
    float ix1, iy1, ix2, iy2, iarea;
    // variables for computing overlap with box j (lower scoring box)
    float xx1, yy1, xx2, yy2;
    float w, h;
    float inter, ovr;

    long num_to_keep = 0;
    for (_i=0; _i < boxes_num; ++_i) {
        i = order_flat[_i];
        if (suppressed_flat[i] == 1) {  // 因为这个已经被之前的bbox给过滤掉了。所以直接pass.            continue;
        }
        //因为这个之前的许多都被过滤掉了,那这个一定是要被保存下来的。
        //比如0号bbox把1,。。。,9都pass掉了,即keep_out_flat[0] = order_flat[0]
        //keep_out_flat[1] = order_flat[10], 这样keep_out_flat最终里面存的就是要保存来的,如
果要取前100个的话,就直接可以取。
        
        keep_out_flat[num_to_keep++] = i;
        // 获得当前的这个bbox的位置信息
        ix1 = boxes_flat[i * boxes_dim];
        iy1 = boxes_flat[i * boxes_dim + 1];
        ix2 = boxes_flat[i * boxes_dim + 2];
        iy2 = boxes_flat[i * boxes_dim + 3];
        iarea = areas_flat[i];
        //从这个开始往后面去看,即算他们的IOU的大小。
        for (_j = _i + 1; _j < boxes_num; ++_j) {
            j = order_flat[_j];
            if (suppressed_flat[j] == 1) {   //同样的道理,如果这个在前面已经被过滤掉了。
                continue;
            }
            xx1 = fmaxf(ix1, boxes_flat[j * boxes_dim]);
            yy1 = fmaxf(iy1, boxes_flat[j * boxes_dim + 1]);
            xx2 = fminf(ix2, boxes_flat[j * boxes_dim + 2]);
            yy2 = fminf(iy2, boxes_flat[j * boxes_dim + 3]);
            w = fmaxf(0.0, xx2 - xx1 + 1);
            h = fmaxf(0.0, yy2 - yy1 + 1);
            inter = w * h;
            ovr = inter / (iarea + areas_flat[j] - inter);
            if (ovr >= nms_overlap_thresh) {  //如果IOU超过了阈值就标记为1.后面就不需要再看>这个了。
                suppressed_flat[j] = 1;
            }
        }
    }

    long *num_out_flat = THLongTensor_data(num_out);
    *num_out_flat = num_to_keep;
    THByteTensor_free(suppressed);
    return 1;
}


综上来看,这里的logic还是非常的清晰的,即用一个suppressed_flat来标记每一个bbox的状态。来记录其是否有没有用,这种方法挺常见的。

然后有几个知识点不是很确定,是我自己理解的。

THArgCheck(THLongTensor_isContiguous(keep_out), 0, "keep_out must be contiguous");这个里面的数字应该是其对应的参数的id, 即它是第0个参数。 THByteTensor这个我的理解是如果用01的话就用这个。即这个张量中的数非0即1. THByteTensor_fill(supressed, 0);这个应该是初始化全为0. 然后用一个指针指向它 unsigned char* suppressed_flat = THByteTensor_data(suppressed);

然后来看gpu的。

gpu最重要的是要理解其内部logic,这一点是写核函数的关键。 比如开多少个block,每个block上面开多少个线程,哪些量需要是shared的,有很多的细节要处理,不然就会常 常 发生越界行为。 这里面的logic是这样的,假设现在有6000个bbox,需要通过NMS, 每个都是有5个信息的。即位置和得分。

那想法是开6000个线程来分别去处理每一个bbox, 但是一个block里面可能开不了那么多的线程,以我的GPU的话,每个block里面最多开1024个线程。那么通过计算,需要6个block才行。 类似于CPU写的代码,这里需要一个mask,它的shape是 (6000, 6),它的意思 是每一行代表一个bbox,其中的6指的是6个block。 比如最前面的6个,里面的数字分别代表着第0号bbox在第0,1,2,3,4,5这几个block里面过滤掉的那些框。而且它存的是从它之后的,即第10号存的就是从11号开始的。

torch和c连接的代码如下

#include <THC/THC.h>
#include <TH/TH.h>
#include <math.h>
#include <stdio.h>

#include "cuda/nms_kernel.h"


extern THCState *state;

int gpu_nms(THLongTensor * keep, THLongTensor* num_out, THCudaTensor * boxes, float nms_overlap_thresh) {
  // boxes has to be sorted
  THArgCheck(THLongTensor_isContiguous(keep), 0, "boxes must be contiguous");
  THArgCheck(THCudaTensor_isContiguous(state, boxes), 2, "boxes must be contiguous");
  // Number of ROIs
  int boxes_num = THCudaTensor_size(state, boxes, 0);   //应该是6000
  printf("boxes_num: %d", boxes_num);
  int boxes_dim = THCudaTensor_size(state, boxes, 1);    //应该是5
  printf("boxes_dim : %d", boxes_dim);
  float* boxes_flat = THCudaTensor_data(state, boxes);  
  const int col_blocks = DIVUP(boxes_num, threadsPerBlock);   //6   blocks的数目。
  THCudaLongTensor * mask = THCudaLongTensor_newWithSize2d(state, boxes_num, col_blocks);   //6000×6的一个张量,但是其实还是个指针,
  unsigned long long* mask_flat = THCudaLongTensor_data(state, mask); 

//执行核函数
 _nms(boxes_num, boxes_flat, mask_flat, nms_overlap_thresh); 


在GPU上算完了之后,下面就是处理那个mask的过程了。

 THLongTensor * mask_cpu = THLongTensor_newWithSize2d(boxes_num, col_blocks);
  THLongTensor_copyCuda(state, mask_cpu, mask);
  THCudaLongTensor_free(state, mask);

  unsigned long long * mask_cpu_flat = THLongTensor_data(mask_cpu);



  THLongTensor * remv_cpu = THLongTensor_newWithSize1d(col_blocks);  //假设长度是6吧,
  unsigned long long* remv_cpu_flat = THLongTensor_data(remv_cpu);
  THLongTensor_fill(remv_cpu, 0);

  long * keep_flat = THLongTensor_data(keep);
  long num_to_keep = 0;


   // 下面的过程就是处理来的那个mask的。
  int i, j;
  for (i = 0; i < boxes_num; i++) {   //相当于是第几个线程来处理第几个bbox
    int nblock = i / threadsPerBlock;  //  block号
    int inblock = i % threadsPerBlock;   // 线程号。


上面涉及到的有从GPU到copy到CPU的函数。然后就是最关键的地方。

   if (!(remv_cpu_flat[nblock] & (1ULL << inblock))) {   //按位与,只要有一个是0就是0.说明当那
个block中的还没有处理的时候,进行处理,这个是对后面的做了个清0操作。
      keep_flat[num_to_keep++] = i;   //第0号就等于i, 并且num_to_keep要加1.
      unsigned long long *p = &mask_cpu_flat[0] + i * col_blocks;   // 指针指向那一行的6个数的>第0个。
      for (j = nblock; j < col_blocks; j++) {   // 只用当前block后面的。对呀,因为前面的可能有>要它算过的了。 
        remv_cpu_flat[j] |= p[j];   //这个长度为6是一直在更新的。
      } 
    }
  }
  
  long * num_out_flat = THLongTensor_data(num_out);
  * num_out_flat = num_to_keep;

  THLongTensor_free(mask_cpu);
  THLongTensor_free(remv_cpu);


这一段写的感觉非常高深,首先是这个if判断这里,其中remv_cpu_flat的长度是6,是记录当前bbox的中的mask中的6个值的。刚才提到这6个mask的值分别代表着那个block中被该bbox给过滤掉的bbox的index。举个例子,比如nblock=0=inblock的时候,就假设是i=0,的时候。 第一次是能够通过if的。然后 remv_cpu_flat就会记录bbox0的mask的6个值。 假设 bbox0把bbox1,2,3,。。。,10全部都给pass掉了,其它的仍然有。 那么此时这个mask的6个值就是[2^1+2^2+...+2^10,0,0,0,0,0],其中用二进表示的时候即 11111111110,那么只有当inblock>=11的时候才会再次走if,因为1ULL<<11 --->1000000000000(11个0). 其中1ULL<<inblock和某个数做按位与的话实际上就是把后面的位数清0。1ULLunsigned long long 类型的1.

感觉这个想法非常厉害。

本来感觉这样算了之后会不会有重复,仔细推算了一下之后,发现其实是不会的。

核函数如何写

#ifdef __cplusplus
extern "C" {
#endif

#include <math.h>
#include <stdio.h>
#include <float.h>
#include "nms_kernel.h"

//这个是计算IOU

__device__ inline float devIoU(float const * const a, float const * const b) {
  float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]);
  float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]);
  float width = fmaxf(right - left + 1, 0.f), height = fmaxf(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);
}

//核函数

__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;//在不同的block中这些数字是不一样的。


  // if (row_start > col_start) return;

  const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);   //是每一行的线程个数吗
,还是shengxia的那些。
  const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

  __shared__ float block_boxes[threadsPerBlock * 5];  //这个是在一个block中的所有的线程都能够看
到的。// 1024×5

 //dev_boxes虽然是6000,5的,但是到这里面之后都当成是数组了,即6000×5的长度。
  if (threadIdx.x < col_size) {
    block_boxes[threadIdx.x * 5 + 0] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0];   //block的两个维度刚好
可以看成是pudiban,.x就是横着,.y就是竖着。
    block_boxes[threadIdx.x * 5 + 1] =
        dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1];   //这个是在dev_boxes中
的位置。
    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();

上面的操作是要把数据全部布到col上面去,然后再按行处理。

 if (threadIdx.x < row_size) {   // 一个block只能处理1024个bbox,所以这个是算当前的block的具体>到某个线程时算的是哪一个bbox.
    const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x;
    const float *cur_box = dev_boxes + cur_box_idx * 5;   //当前的box的第一个地址。这里的"+"相>当于是移位。 //用这个指针来得到当前bbox这5个数的首地址。
    //现在就想像在某一个block里面的某个线程正在做什么?

    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;   //t的值就变了。
      }
    }
    const int col_blocks = DIVUP(n_boxes, threadsPerBlock);   //blocks的数目 。
    dev_mask[cur_box_idx * col_blocks + col_start] = t;   //dev_mask就是记录的。   //不改就是0,改了至少是1
  }
}


//最主要的就是这个函数。
// 6000, boxes,
void _nms(int boxes_num, float * boxes_dev,
          unsigned long long * mask_dev, float nms_overlap_thresh) {

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

#ifdef __cplusplus
}
#endif
                                                                             102,1  

打赏,谢谢~~

取消

感谢您的支持,我会继续努力的!

扫码支持
扫码打赏,多谢支持~

打开微信扫一扫,即可进行扫码打赏哦