Cython学习小结及代码实践(NMS非极大值抑制的Cython实现)

在Jupyter中使用Cython

• 在cython代码之前的单元格中添加并运行%load_ext Cython命令
• 在cython代码的单元格第一行添加%%cython命令
• %%cython命令的几个参数: 1. -a: 表示编译后显示代码注释. 2. -cplus表示使用c++编译. (其余的待补充)

用Cython改写python代码的要点

• Cython加速python代码的关键在于指定变量类型, 用内存视图(memoryview)进行索引, 将代码改写为原子操作(利用循环的方式对复合过程进行拆解).
• 在纯python代码中所有类型都是一个对象(底层是c语言的结构体), 因此即便是int类型这些看上去像数值类型的类型, 其底层的存储方式和运算方式也并非如同表面上那么简单. 不同于c语言整数的存储和运算方式, 类似于在python中自定义一个类然后利用运算符重载的方式让两个类看起来可以运算.
• 由于上述原因, for和while循环体常常是python代码低效运行的地方(因为要重复对python对象而不是c的数据类型进行运算). 但是如果是对c语言的数据类型进行循环运算则不受影响, 反而更高效(因为c语言的数据类型存储时有固定的大小, 方便分配内存, 也因此可以通过计算指针位置的方式高效的索引).
• 对应于上述的原因可知:

1. 指定变量类型的作用是将变量申明为明确的c类型(当然每个用c实现的python中, 每个对象的底层也都是一个复杂的c结构体, cython中也可以申明像numpy.ndarray这样复杂的类型, 但是这通常不会带来加速, 因为本质上还是一个复杂的数据类型);
2. 内存视图的作用是对类似c语言中数组的内存区域进行高效索引(numpy数组实际上已经可以做到很高效的访问数组了, 但据cython官网上所说, numpy的一些实现会对内存进行复制, 这会带来性能上的损耗, 也是cython可以优化的比numpy快的原因, 当然负面影响是可能会影响安全性, 不合法的操作可能会数据造成破坏).
3. 将代码改写为原子操作的作用是避免和python进行太多交互(cython内部会将一个python对象转换成c数据, 还会附加一系列验证等代码, 需要时还会将c数据结果转换成python对象返回, 这些都会降低运算速度, 特别是出在循环体中), 直接对简单的c数值类型进行运算, 这通常需要将一个python函数利用循环展开.

代码实践: 用Cython改写nms代码

• 简介: nms(非极大值抑制)是目标检测中的一种常用操作, 其目的是过滤掉和高分框重叠度较高的框.
• 说明: 单个detection box的数据结构为[x1, y1, x2, y2], 分别为左上角和右下角的坐标.
• 说明: 下面的代码都在jupyter notebook中运行和测试
• 伪代码如下:

Input: B = {b1, ..., bN}, S = {s1, ..., sN}, Nt
        B是一系列初始检测框
        S是对应的检测框评分
        Nt是NMS阈值

begin:
    D <-- {}
    while B != empty do:
        m <-- argmax S
        M <-- bm
        D <-- D | M; B <-- B - M
        for bi in B do:
            if iou(M, bi) >= Nt then:
                B <-- B - bi; S <-- S - si
            end
        end
    end
    return D, S
end

1. 网上的实现(纯python实现)

• 来源: https://www.cnblogs.com/king-lps/p/9031568.html

• 使用numpy的实现作为基准.

代码部分

def py_cpu_nms(dets, thresh):
    # dets:(m,5)  thresh:scaler
    
    x1 = dets[:,0]
    y1 = dets[:,1]
    x2 = dets[:,2]
    y2 = dets[:,3]
    
    areas = (y2-y1+1) * (x2-x1+1)
    scores = dets[:,4]
    keep = []
    
    index = scores.argsort()[::-1]
    
    while index.size >0:

        i = index[0]       # every time the first is the biggst, and add it directly
        keep.append(i)
        
        x11 = np.maximum(x1[i], x1[index[1:]])    # calculate the points of overlap 
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])
        
        w = np.maximum(0, x22-x11+1)    # the weights of overlap
        h = np.maximum(0, y22-y11+1)    # the height of overlap
       
        overlaps = w*h
        
        ious = overlaps / (areas[i]+areas[index[1:]] - overlaps)
        
        idx = np.where(ious<=thresh)[0]
        
        index = index[idx+1]   # because index start from 1
        
    return keep

测试代码部分

为了增加区分度, 模拟对100000个检测框进行nms.

%%timeit

np.random.seed(10)
dets = np.sort(np.random.uniform(0, 2000, size=(100000, 5)), axis=-1).astype(np.float32)
dets[:, 4] = np.random.uniform(0, 1, size=100000).astype(np.float32)
iou_thresh = 0.7
py_cpu_nms(dets, iou_thresh)

测试结果

4.99 s ± 53.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

2. 期望达到的目标(pytorch的实现)

• pytorch的torchvision库有nms的实现, 其cpu版本的实现基于c++, 我们以此作为优化的目标.
• 接口和上述的实现稍有不同, boxes和scores是分别输入
• 底层的实现也稍有不同, 体现在对area面积的计算上, 其最后剩余的框也更多, 因此与基准相比, 其实际速度比测试中体现得还要更快一点.

代码部分

from torchvision.ops import nms
import torch

测试部分

np.random.seed(10)
dets = np.sort(np.random.uniform(0, 2000, size=(100000, 5)), axis=-1).astype(np.float32)
dets[:, 4] = np.random.uniform(0, 1, size=100000).astype(np.float32)
iou_thresh = 0.7
nms(torch.tensor(dets[:, :4]), torch.tensor(dets[:, 4]), iou_thresh)

测试结果

2.24 s ± 35 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

3. 网上的Cython实现

• 来源同上
• 原文代码中计算iou面积处有错误, 导致会迅速过滤掉大量检测框, 因此原文的测试显的很快. 下面的代码对相应的地方进行了修正.

代码部分

%%cython -a --cplus
import numpy as np
cimport numpy as np
 
 
cdef inline np.float32_t max(np.float32_t a, np.float32_t b):
    return a if a >= b else b

cdef inline np.float32_t min(np.float32_t a, np.float32_t b):
    return a if a <= b else b

def py_cpu_nms_cython(np.ndarray[np.float32_t,ndim=2] dets, np.float thresh):
    # dets:(m,5)  thresh:scaler
    
    cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:,0]
    cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:,1]
    cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:,2]
    cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:,3]
    cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4]
    
    cdef np.ndarray[np.float32_t, ndim=1] areas = (y2-y1+1) * (x2-x1+1)
    cdef np.ndarray[np.int_t, ndim=1]  index = scores.argsort()[::-1]    # can be rewriten
    keep = []
    
    cdef int ndets = dets.shape[0]
    cdef np.ndarray[np.int_t, ndim=1] suppressed = np.zeros(ndets, dtype=np.int)
    
    cdef int _i, _j
    
    cdef int i, j
    
    cdef np.float32_t ix1, iy1, ix2, iy2, iarea
    cdef np.float32_t w, h
    cdef np.float32_t overlap, ious
    
    j=0
    
    for _i in range(ndets):
        i = index[_i]
        
        if suppressed[i] == 1:
            continue
        keep.append(i)
        
        ix1 = x1[i]
        iy1 = y1[i]
        ix2 = x2[i]
        iy2 = y2[i]
        
        iarea = areas[i]
        
        for _j in range(_i+1, ndets):
            j = index[_j]
            if suppressed[j] == 1:
                continue
            xx1 = max(ix1, x1[j])
            yy1 = max(iy1, y1[j])
            xx2 = min(ix2, x2[j])
            yy2 = min(iy2, y2[j])
    
            w = max(0.0, xx2-xx1+1)
            h = max(0.0, yy2-yy1+1)
            
            overlap = w*h 
            ious = overlap / (iarea + areas[j] - overlap)
            if ious>thresh:
                suppressed[j] = 1
    
    return keep

测试部分

%%timeit

np.random.seed(10)
dets = np.sort(np.random.uniform(0, 2000, size=(100000, 5)), axis=-1).astype(np.float32)
dets[:, 4] = np.random.uniform(0, 1, size=100000).astype(np.float32)
iou_thresh = 0.7
py_cpu_nms_cython(dets, iou_thresh)

测试结果

16.6 s ± 315 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

测试结果说明

• 可以看到该实现是比基准(也就是纯python版)还要慢2倍的.
• 通常来说python代码直接复制为cython代码, 其速度是会提升10个百分点的.
• 可以从cython生成的注释报告中看出以下原因, 这也是我们在改写cython代码时需要注意的:
  
  
  

代码分析:

• 图中黄色表示代码与python的交互(即cython会将python代码和c代码相互转换, 点击代码行左侧的"+"可以查看详细情况), 颜色越深表示交互越多.
• 可以看到将dets变量申明为np.ndarray(11到25行)引入了大量的交互代码(颜色深), 通常在循环外的交互不会对最终性能造成太多影响, 因为只运行有限次, 申明为ndarray的而不是memoryview的好处是, ndarray仍可以使用np的函数(如21行处的argsort), 但官方推荐numpy用户采用memoryview的形式, 因为更通用更高效.
• 真正会对性能造成致命影响的是在内层循环处与python交互(55行到65行), 其交互的原因是没有申明xx1, yy1等变量的类型, 导致在这些变量的赋值会从c数据转换成python对象, 增加了许多类型转换和检查的代码, 60和61行的运算也变成了python对象间的运算. 如果这些代码出现在循环中, 特别是内层循环, 将极大的影响速度.

代码修改

修改的地方包括以下几点:

• 将参数thresh的申明从np.float改为np.float32_t. 原因是np.float是python类型, 而后者是c类型. 改为float也可以, 因为在cython中直接用int和float等申明的是c类型.
• 添加xx1, xx2, yy1, yy2的申明类型
• 对函数进行装饰, 表明不对index, 是否除0等进行检查.
• cython报告中显示循环内部基本没有交互代码了

修改的测试结果

• 可以看到终于比基准有一点进步. 说明numpy确实优化的很好

4.65 s ± 10.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

4. 向目标迈进

• 上面的代码中还有哪些地方可以改进呢? 一般的思路都是优化循环体中的索引部分, 或者根据cython报告查看哪些地方需要改进, 或者更具体的是用line_profiler查看每行cython代码的运行次数和时间, 此处就不详细记录了.
• 下面的代码结合了上述所有提到的可优化的点, 由于是自己写的, 所以和上述他人的代码变量名不一致, 但思路大体一致. 不一样的地方是使用cython的memoryview来优化索引, 以及在计算相交面积前进行判断, 避免不必要的计算. 另外不一样的地方是对面积的计算, 我计算面积的方法和pytorch的实现方法一致, 因此和Pytorch的结果是一样的.
• 没有将函数参数中的boxes和scores直接申明为内存视图的原因是函数中使用了numpy的函数, 更重要的原因在后面会提到.

代码部分

%%cython -a --cplus
import numpy as np
cimport numpy as np
cimport cython


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def nms_cython(np.ndarray[np.float32_t, ndim=2] boxes, np.ndarray[np.float32_t, ndim=1] scores, float iou_threshold):
    cdef Py_ssize_t N = boxes.shape[0]
    cdef float [:, :] boxes_view = boxes
    result = []
    cdef float [:] area = np.zeros(N, dtype=np.float32)
    
    cdef Py_ssize_t _i_area = 0
    for _i_area in range(N):
         area[_i_area] = (boxes_view[_i_area, 2] - boxes_view[_i_area, 0]) *  (boxes_view[_i_area, 3] - boxes_view[_i_area, 1])
        
    cdef:
        Py_ssize_t index_current, index_ori, index_compared_ori
        Py_ssize_t [:] index_argsort = np.asarray(scores).argsort()[::-1].astype(np.int64)
        Py_ssize_t [:] index_ignore = np.zeros(N, dtype=np.int64)
        float _x11, _x21, _x12, _x22, _y11, _y21, _y12, _y22
        float area_intersection, area_total, iou
    
    for index_current in range(N):
        if index_ignore[index_current] == 1:
            continue      
            
        index_ori = index_argsort[index_current]
        result.append(index_ori)
        
        _x11 = boxes_view[index_ori, 0]
        _y11 = boxes_view[index_ori, 1]
        _x12 = boxes_view[index_ori, 2]
        _y12 = boxes_view[index_ori, 3]
        for index_compared_box in range(index_current + 1, N):
            if index_ignore[index_compared_box] == 1:
                continue   
            index_compared_ori = index_argsort[index_compared_box]
            
            _x21 = boxes_view[index_compared_ori, 0]
            _y21 = boxes_view[index_compared_ori, 1]
            _x22 = boxes_view[index_compared_ori, 2]
            _y22 = boxes_view[index_compared_ori, 3]
            
            if  min(_y12, _y22) - max(_y11, _y21) < 0  or (min(_x12, _x22) - max(_x11, _x21)) < 0:
                continue
            
            area_intersection = (min(_y12, _y22) - max(_y11, _y21)) * (min(_x12, _x22) - max(_x11, _x21))
            area_total = area[index_ori] + area[index_compared_ori]
            iou = area_intersection / (area_total - area_intersection)
            if iou > iou_threshold:
                index_ignore[index_compared_box] = 1
    return result

测试结果

测试用例和前面相同.

• 可以看到本次的优化已经比numpy快很多了. 但离pytorch的实现还有距离.

2.89 s ± 10.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

5. 最终优化的结果

• 上面的代码中一直在回避的一个问题是, 在for循环中, 我们的策略是碰到标记为已经抑制的框就continue略过此次循环. 当候选框数量很大的时候, 这会影响速度(因为即便什么都不做但此次遍历还是会耗费时间). 在纯python的实现中, 已经标记为抑制的框下次就不会再遍历了, 这是通过布尔索引和重新赋值来实现的. 但是cython里的内存视图不具有这样的功能.
• 如果有像python中的集合或者动态列表则很容易解决这个问题, 但在cython最好不要使用python对象. 这就需要用到c++中的list数据结构, 它的底层实现是一个双向链表, 删除操作的时间复杂度为1, cython使用c++的数据类型不会带来太多额外的开销.
• 以下最终优化的代码, 已经远没有原来的python代码那样清晰了, 这也是优化带来的弊端. 由于使用c++的数据结构更多要注意的是实现的细节, 代码中给出我的一个实现供大家参考, 更多的用法需要去官网了解和对c++的学习.
• 其他优化的点是, 使用ascontiguousarray传递一个连续内存的副本给cython的memoryview, 经我的实践是很有用的. 其内在的原因, 我个人的理解是和cpu工作的原理有关, 虽然RAM可以做随机访问, 但cpu在访存时并非一次只读一个地址的数据, 而是会将一块邻近的代码一次加载到缓存中, 如果下次访问的代码在缓存中命中, 则直接从缓存而不是内存中读取, 这会带来速度上的提升. (所以访问顺序可能也有影响?读数时尽量按x1, y1, x2, y2的顺序读?这点我没有验证)
• 这几个优化的点都是琐碎的, 且需要在实践中验证的, 通常带来的收益极为有限. 还是将代码表述的更清晰更加重要, 需要优化的时候再进行优化

代码部分

%%cython -a --cplus

import numpy as np
cimport numpy as np
cimport cython
from libcpp.list cimport list as cpplist
from cython.operator cimport dereference as deref, preincrement as inc, predecrement as dec

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def nms_cython_optim(np.ndarray[np.float32_t, ndim=2] boxes, np.ndarray[np.float32_t, ndim=1] scores, float iou_threshold):
    cdef Py_ssize_t N = boxes.shape[0]
    cdef float [:, ::1] boxes_view = np.ascontiguousarray(boxes)
    result = []
    
    cdef float [:] area = np.zeros(N, dtype=np.float32)
    cdef Py_ssize_t _i_area = 0
    for _i_area in range(N):
            area[_i_area] = (boxes_view[_i_area, 2] - boxes_view[_i_area, 0]) *  (boxes_view[_i_area, 3] - boxes_view[_i_area, 1])
        
    index_argsort_numpy = np.ascontiguousarray(scores.argsort()[::-1].astype(np.int64))
    
    cdef:
        Py_ssize_t index_current, index_ori, index_compared_ori
        Py_ssize_t [::1] index_argsort = index_argsort_numpy
        Py_ssize_t [::1] index_ignore = np.zeros(N, dtype=np.int64)
        float _x11, _x21, _x12, _x22, _y11, _y21, _y12, _y22
        float area_intersection, area_total, iou, _area_current
        cpplist[Py_ssize_t] keep_index = index_argsort_numpy
        cpplist[Py_ssize_t].iterator current_index_iter
    
    while keep_index.size() > 0:
        current_index_iter = keep_index.begin()
        index_ori = deref(current_index_iter)
        result.append(index_ori)
        current_index_iter = keep_index.erase(current_index_iter)

        _x11 = boxes_view[index_ori, 0]
        _y11 = boxes_view[index_ori, 1]
        _x12 = boxes_view[index_ori, 2]
        _y12 = boxes_view[index_ori, 3]
        
        _area_current = area[index_ori]
        while current_index_iter != keep_index.end():
            index_compared_ori = deref(current_index_iter)
            
            _x21 = boxes_view[index_compared_ori, 0]
            _x22 = boxes_view[index_compared_ori, 2]
    
            if (min(_x12, _x22) - max(_x11, _x21)) < 0:
                inc(current_index_iter)
                continue
                
            _y21 = boxes_view[index_compared_ori, 1]
            _y22 = boxes_view[index_compared_ori, 3]
            
            if  min(_y12, _y22) - max(_y11, _y21) < 0:
                inc(current_index_iter)
                continue
            
            area_intersection = (min(_y12, _y22) - max(_y11, _y21)) * (min(_x12, _x22) - max(_x11, _x21))
            area_total = _area_current + area[index_compared_ori]
            iou = area_intersection / (area_total - area_intersection)
            if iou > iou_threshold:
                current_index_iter = keep_index.erase(current_index_iter)
            else:
                inc(current_index_iter)
    return result

测试结果

• 已经和pytorch的结果相当了!
• 虽然提升的幅度没有一开始那么大, 但是越到后面优化起来就越困难. 可以设想再要提升, 只能考虑并发编程了, 这又是一个复杂的问题, 就先告一段落了.

2.22 s ± 12.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

其他的说明

• 由于cython是边学边练的, 接下来还记录一些学习过程中的测试.
• memoryview是否可以使用numpy函数? 答: 不可以, 但是可以通过np.asarray(memoryview)来使用, numpy的ndarray和cython的memoryview共用一个buffer, 谁修改都可以, 所以可以要高效索引时使用view, 要用numpy函数时用对应的ndarray.
• memoryview可不可以用bool索引, 花式索引等等? 不可以.
• cython中的c对象可不可以申明相同名称的变量? 不可以, 所以上面不得以使用c++的list结构, 因此不可以for循环中(其他地方也不可以)对同一个变量申明不同长度的内在视图.
• 如何使用line_profiler查看cython代码的profile?

1 .cython代码前面的单元格需要添加并运行%load_ext line_profiler;
2. cython代码的单元格中需要注释(会影响运算, 分析完后需要删掉重新编译):

# cython: linetrace=True
# cython: binding=True
# distutils: define_macros=CYTHON_TRACE_NOGIL=1

3. cython中改写的函数由def改为cpdef
4. 添加一个哨兵函数(任意形式):
   (注: 第3第4点据说是因为不这么写line_profiler会有bug, 不显示结果, 未验证这一说法)

def nms_cython_optim_sentinel():
    pass

5. 用于测试的单元格中输入:

%lprun -f nms_cython_optim nms_cython_optim(test_boxes, test_scores, iou_thresh)