在Python中用cupy实现IoU（交并比）和NMS（非极大值抑制）的GPU加速

1. 前言

IoU（交并比）和NMS（非极大值抑制）的计算在目标检测任务中可以说是必不可少的，但是当需要计算的bounding box的数量级很大的时候，cpu就吃不消了。例如在对Faster RCNN的RPN网络进行训练的时候，需要我们计算近20000个anchor bbox与ground-truth bbox之间的IoU；另外，我们还需要从RPN输出的proposal RoI中选出最后RCNN的训练样本，这一过程也需要对10000+的bbox进行NMS处理。这时就需要我们将其中的数组运算放入GPU中来执行，利用其并行处理能力来实现加速。

有些时候我们可以将这两个过程纳入深度学习框架内来实现GPU加速，但很多时候用这些框架（比如Tensorflow）提供的数组类型实现起来并不方便，因此本文提供的方法主要是通过CuPy库提供的API来实现IoU和NMS的CUDA加速。（顺带一提，CuPy此前是深度学习框架Chainer下基于cuda的底层计算引擎，现在已经成为一个独立的库，和它类似的库还有PyCuda，而之所以选择CuPy其实只是看到它在GitHub上的star比较多而已Orz…不过稍微看了一下，感觉两者的基本函数语法应该大同小异…

2. CuPy中的ElementwiseKernel类

首先我们需要把IoU的计算过程封装成一个核函数，这里的核（Kernel）是CUDA编程中的一个概念，我了解也不多，但可以简单理解为在GPU中执行运算的函数。CuPy为用户提供了多种自定义核函数的方法，具体可以看官方的说明文档，这里我们用到的是CuPy的ElementwiseKernel类。顾名思义，通过这个类定义的核执行的是element-wise的运算，下面是自己简单总结的几个注意点：

• 每个element的运算都定义在一个“循环体”（loop body）中，不过每次“循环”之间其实是并行的。
• “循环”的次数由输出数组的形状决定，而输出数组的形状可以交由ElementwiseKernel类自身根据输入数组的形状推断（支持boardcasting），也可以由用户自己指定。
• 但如果需要在循环体使用到数组索引，则必须自己指定输出数组的形状。
• 数组在被输入核中后会被flatten成一维的数组，所以索引的时候也要按照flatten之后的位置来索引。
• 输入核中的变量都是只读的，不能进行赋值。

关于这个类以及其他kernel类的具体语法和一些简单的示例也可以从上面给出的官方说明文档中获取，这里不作赘述。

3. 实现IoU的加速

好了，知道了怎么用CuPy自定义element-wise kernel，我们就可以用它写在GPU上跑的IoU函数啦：

import cupy as cp
import numpy as np

bbox_iou = cp.ElementwiseKernel(
'raw T bbox_a, raw T bbox_b, int32 num_b',
'float32 iou',
'''
//int num_b = sizeof(bbox_b) / (sizeof(bbox_b[0])*8);
int i_a = i / num_b; //bbox_a数组中第？个框的索引值
int i_b = i % num_b; //bbox_b数组中第？个框的索引值

//bbox_a中第i_a个框的面积
T width_a = bbox_a[4*i_a + 2] - bbox_a[4*i_a + 0];
T height_a = bbox_a[4*i_a + 3] - bbox_a[4*i_a + 1];
T area_a = width_a * height_a;

//bbox_b中第i_b个框的面积
T width_b = bbox_b[4*i_b + 2] - bbox_b[4*i_b + 0];
T height_b = bbox_b[4*i_b + 3] - bbox_b[4*i_b + 1];
T area_b = width_b * height_b;

//相交区域的的面积
T left = max(bbox_a[4*i_a + 0], bbox_b[4*i_b + 0]);
T top = max(bbox_a[4*i_a + 1], bbox_b[4*i_b + 1]);
T right = min(bbox_a[4*i_a + 2], bbox_b[4*i_b + 2]);
T bottom = min(bbox_a[4*i_a + 3], bbox_b[4*i_b + 3]);
T width_inter = max(right - left, static_cast(0));
T height_inter = max(bottom - top, static_cast(0));
T area_inter = width_inter * height_inter;

iou = area_inter / (area_a + area_b - area_inter + 1e-6);
''',
'IoU')

代码中提供了必要的注释，怎么样，挺简单的对吧？函数实现的是包含N个边界框的数组bbox_a和包含K个边界框的数组bbox_b的IoU计算，输出是一个形状为(N,K)的二维数组，其实第[n, k]个元素表示的是bbox_a中第n个框和bbox_b中第k个框的IoU值。

下面我们用CuPy的array定义几个简单的边界框，送进去试一试：

a = cp.array([[2, 1, 5, 6],[6, 7, 9, 10], [1, 2, 3, 4]], dtype=cp.float32).reshape(3, 1, 4)
b = cp.array([[2, 1, 5, 6],[6, 7.3, 9, 10.1]], dtype=cp.float32).reshape(1, 2, 4)
y = cp.empty((3, 2), dtype=cp.float32)  #需指定输出数组的形状
out = bbox_iou(a, b, 2, y)  #注意函数输入除了两个bbox数组外，还需要输入bbox_b中框的数目，以及指定的y
print(out.shape)
print(out)

输出情况如下：

(3, 2)

[[ 0.99999994  0.        ]
 [ 0.          0.87096739]
 [ 0.11764705  0.        ]]

4. 实现NMS的加速

非极大值抑制（NMS）本质是排除掉邻域范围内的非极大值，关于NMS的原理和算法流程有很多优秀的博文可以参考，这里也不做赘述了。对NMS的CUDA加速的实现需要一丢丢的逻辑思维，但是也不难。我们只需要知道NMS算法在CPU中运算的瓶颈其实还是在于数组中当前的bbox和剩余的bbox的IoU计算过程，于是我们只需要在核函数中完成这一计算即可，所以NMS的加速本质上还是对IoU运算的加速。与上一节唯一的不同是两个不同数组之间IoU的计算变为了单一数组中每个bbox之间IoU的计算。

假设输出的数组包含N个边界框，那么我们的输出目标就是输出一个形状为(N, N)的数组。但是我们知道，第i个bbox和第j个bbox的IoU与第j个bbox和第i个bbox的IoU是一样的，因此输出的数组其实是个对称矩阵（有点像协方差矩阵呢），另外，某个bbox和自身计算IoU是没有意义的，所以这个矩阵的下三角部分我们其实是不需要的，我们可以置为零。得到这个矩阵后，我们就可以把IoU值大于阈值的置为1，小于阈值的置为0。后续我们可以根据这个二值矩阵来决定这些边界框的去留，但现在我们先以输出这个矩阵为目标写出kernel函数吧：

import cupy as cp
import numpy as np

nms_gpu = cp.ElementwiseKernel(
'raw T bbox, int32 num_box, T threshold',
#'float32 iou',
'uint8 remove_cdd',
'''
int i_x = i % num_box;
int i_y = i / num_box;

// Sinve IoU(a, b) == IoU(b, a) and IoU(a, a) is meaningless, we don't need to 
// calculate the lower triangular part and the diagnal part.
if (i_y < i_x)
{
    T width_x = bbox[4*i_x + 2] - bbox[4*i_x + 0];
    T height_x = bbox[4*i_x + 3] - bbox[4*i_x + 1];
    T area_x = width_x * height_x;
    
    T width_y = bbox[4*i_y + 2] - bbox[4*i_y + 0];
    T height_y = bbox[4*i_y + 3] - bbox[4*i_y + 1];
    T area_y = width_y * height_y;
    
    T left = max(bbox[4*i_x + 0], bbox[4*i_y + 0]);
    T top = max(bbox[4*i_x + 1], bbox[4*i_y + 1]);
    T right = min(bbox[4*i_x + 2], bbox[4*i_y + 2]);
    T bottom = min(bbox[4*i_x + 3], bbox[4*i_y + 3]);
    T width_inter = max(right - left, static_cast(0));
    T height_inter = max(bottom - top, static_cast(0));
    T area_inter = width_inter * height_inter;
    
    float iou = area_inter / (area_x + area_y - area_inter + 1e-6); 
    remove_cdd = (iou > threshold); //被标记的1的位置的列索引其实是在后续可能会被移除的候补框
}

''',
'nms')

可以看到代码内容其实跟上一节大同小异，只是多了个if语句以避免重复计算的IoU和对角线上的IoU。我们再随便定义几个bbox，输入函数中试一下：

x = cp.array([[2, 1, 5, 6],[2.6, 1.1, 5, 6], [1, 2, 3, 4], [2.9, 1.1, 5,6], [3, 1.5 ,5.2, 6]], dtype=cp.float32)
y = cp.zeros((num_bbox, num_bbox), dtype=cp.uint8) 
rmv = nms_gpu(x, num_bbox, 0.6, y)
print(rmv)

输出如下：

[[0 1 0 1 0]
 [0 0 0 1 1]
 [0 0 0 0 0]
 [0 0 0 0 1]
 [0 0 0 0 0]]

可以看到，输出矩阵rmv的[i, j]位置上的数值其实反映了第i个bbox和第j个bbox的IoU是否大于阈值。

接下来需要明确，我们最终希望通过NMS得到是一个与bbox长度一致的掩膜数组，其元素标记了对应位置的bbox是否应该被移除。下面解释下实现的过程，大家也可以直接拉下去看代码，也不难懂。为了实现输出的目标，我们可以先定义一个长度为N初始值为0的一维数组mask，然后根据mask中的值对rmv由上到下进行逐行的判断（注意输入的bbox数组其实已经根据分类的probability进行了从大到小的排序，所以矩阵处理的优先顺序自然也是由上到下），再通过或运算来更新mask。比方说进行到rmv矩阵的第i行时，如果当前mask的第i个数为0， 那么说明第i个bbox是需要保留的，那么矩阵这一行的值与mask作或运算来更新mask；若当前mask的第i个数为1，说明第i个bbox已经确定了要被移除，则跳过这一行，mask不作更新。这样，当遍历完整个矩阵，最终得到的就是我们需要的掩膜了。可以说这一段代码才是真正的“非极大值抑制”的过程。

下面给出代码，但要注意的是与上面的核函数不同，这段代码其实是在cpu上执行的，我也尝试过用CuPy的数组来写，不过事实证明这种顺序执行的逻辑运算还是CPU比较适合。当然用Cython来写运行速度应该会更快，不过我不会Cython，有兴趣的同学可以自己试一下。

#接上文的代码
rmv_host = rmv.get() #先把cupy array数组转换为numpy array，也就是把数据从device端移回host端
mask = np.zeros(num_bbox, dtype=np.uint8)
for i in range(0, num_bbox-1):
    if not mask[i]:
        mask |= rmv_host[i] # mask和rmv_host的第i行做或运算
print(mask)
idx = np.where(mask==0)[0] #输出最后保留下来的bbox的索引
print(idx)

以下是输出：

[0 1 0 1 0]
[0 2 4]

破费！第一行输出是我们的二值mask，第二行是最后保留下来的bbox的索引。

这段代码实际上参考了Faster RCNN作者的源码，不过原作者的思路实在是牛匹得多，把每64个bbox分配一个GPU的block中，然后用64位ULL型的二进制数来表示二值mask，再通过位或来更新mask…实在是巧妙。关于Faster RCNN的NMS源码部分的分析，可以去看看这篇博文，内容非常详实到位：目标检测中NMS的GPU实现（来自于Faster R-CNN中的nms_kernel.cu文件）。

最后我们来验证一下NMS通过GPU加速的效果，先基于numpy数组定义一个在cpu下运行的nms函数（实际上这段只定义了计算IoU和比较阈值的部分）：

def nms_cpu(bbox, num_box, threshold):
    remove = np.zeros((num_box, num_box), dtype=np.uint8)
    for i in range(num_box):
        for j in range(num_box):  
            if i < j:
                width_x = bbox[j, 2] - bbox[j, 0]
                height_x = bbox[j, 3] - bbox[j, 1]
                area_x = width_x * height_x
                
                width_y = bbox[i, 2] - bbox[i, 0]
                height_y = bbox[i, 3] - bbox[i, 1]
                area_y = width_y * height_y
                
                left = max(bbox[j, 0], bbox[i, 0])
                top = max(bbox[j, 1], bbox[i, 1])
                right = min(bbox[j, 2], bbox[i, 2])
                bottom = min(bbox[j, 3], bbox[i, 3])
                width_inter = max(right - left, 0.)
                height_inter = max(bottom - top, 0.)
                area_inter = width_inter * height_inter
                
                iou = area_inter / (area_x + area_y - area_inter + 1e-6) 
                remove[i, j] = (iou > threshold)
    return remove

然后我们分别使用Numpy array和CuPy array各定义一个包含5000个bbox的数组（为了方便数组元素全设为1了），然后分别送入CPU的函数和GPU的函数中，看看效果如何：

# 接上文代码，由于后续产生mask的程序是一样的，我们只比较计算iou矩阵的部分
# 顺带一提，电脑是win10系统，gpu是1080TI，cpu是i7-8700
import time

num_bbox = 5000

s = time.time()
a = np.ones((num_bbox, 4), dtype=np.float32) 
rmv_2 = nms_cpu(a, num_bbox, 0.5)
e = time.time()
print('加速前：', e - s, 's')

s = time.time()
b = cp.ones((num_bbox, 4), dtype=cp.float32) 
z = cp.zeros((num_bbox, num_bbox), dtype=cp.uint8) 
rmv_1 = nms(b, num_bbox, 0.5, z)
e = time.time()
print('加速后：', e - s, 's')

输出结果：

加速前： 96.38717579841614 s
加速后： 0.006020069122314453 s

卧槽，直接起飞有没有！？

5. 后记

最近在其他博文中学到了不少，然后大多数博主对于我留下的问题和评论都进行了耐心的解答，非常感谢他们，于是也想着把自己的一些进展和经验分享一下，虽然是一些没啥技术含量的东西，但也希望能多少帮助到大家。另外，我自己其实也只是初学者——无论是在Python编程还是在深度学习方面，所以如果文章中有疏漏和错误的地方，欢迎大家提出和指正~