简单的numba + CUDA 实测

简单的numba + CUDA 实测

起因

一时兴起，是我太闲了吧。

最近需要对一个4k图像进行单个像素级别的处理，由于用python用得人有点懒，直接上python在所有像素上循环一遍。每个像素做的工作其实很简单，就是判断一下这个像素是否符合某一准则，如果符合就将这个像素mask上。噼里啪啦写好一个脚本，一运行居然等了很久，都没有结果。一开还以为是不是哪里写错了，进入了无限循环什么的，但是最后发现其实执行效率就是那么低。我做了一个实例，在一个3008x4112像素的图像上进行简单的分类。（实例中的两个分类功能都可以用cv2直接实现，这里仅作实例进行测试使用。）

from __future__ import print_function

import cv2
import math
import numpy as np
from numpy.linalg import norm
import time

H = 3008
W = 4112

class Validator(object):
    def __init__(self):
        pass

    def is_valid(self, x, y):
        return False

class RadiusValidator(Validator):
    def __init__(self, center, R, width):
        super(RadiusValidator, self).__init__()

        self.R = R
        self.center = center # A two element NumPy array. [x, y].
        self.width = width

        if ( self.width <= 0 ):
            raise Exception("self.width wrong. self.width = {}".format(self.width))
    
    # Overide parent's function.
    def is_valid(self, x, y):
        x = x - self.center[0]
        y = y - self.center[1]

        r = math.sqrt( x * x + y * y )

        if ( r >= self.R - self.width and r <= self.R + self.width ):
            return True
        else:
            return False

class PolarLineSegmentValidator(Validator):
    def __init__(self, center, theta, length, width):
        super(PolarLineSegmentValidator, self).__init__()

        self.center = center # A two element NumPy array. [x, y].
        self.theta  = theta
        self.length = length
        self.width  = width

        self.endP    = np.zeros((2,), dtype=np.float32)
        self.endP[0] = self.length * math.cos( self.theta )
        self.endP[1] = self.length * math.sin( self.theta )

        self.dir = self.endP / self.length

    def is_valid(self, x, y):
        # Transform the coordinate to the local frame.
        v = np.zeros((2, ), dtype=np.float32)

        v[0] = x - self.center[0]
        v[1] = y - self.center[1]

        # Projection.
        proj = self.dir.dot( v )

        if ( proj < 0 ):
            return False
        elif ( proj > self.length ):
            return False

        # Othogonal vector.
        oth = v - proj*self.dir

        # Distance.
        d = norm( oth )

        if ( d <= self.width ):
            return True
        else:
            return False

if __name__ == "__main__":
    print("List the indices for AEAG brightness evaluation.")

    # Create a blank image with all zeros.
    img = np.zeros((H, W), dtype=np.uint8)

    imgCenter = np.array([ W/2, H/2 ], dtype=np.float32)

    # Create the validators.
    vList = []
    vList.append( RadiusValidator( imgCenter, min( W/2, H/2 ) * 0.75, 20 ) )
    vList.append( PolarLineSegmentValidator( imgCenter,  math.pi/2, min( W/2, H/2 ) * 0.75, 10 ) )
    vList.append( PolarLineSegmentValidator( imgCenter,  math.pi/4, min( W/2, H/2 ) * 0.75, 10 ) )
    vList.append( PolarLineSegmentValidator( imgCenter,          0, min( W/2, H/2 ) * 0.75, 10 ) )
    vList.append( PolarLineSegmentValidator( imgCenter, -math.pi/4, min( W/2, H/2 ) * 0.75, 10 ) )
    vList.append( PolarLineSegmentValidator( imgCenter, -math.pi/2, min( W/2, H/2 ) * 0.75, 10 ) )

    # Record the starting time.
    start = time.time()

    for i in range(H):
        for j in range(W):
            flag = False
            for v in vList:
                flag = flag or v.is_valid( j, i )
            
            if ( True == flag ):
                img[i, j] = 255

        if ( 0 == i % 100 ):
            print("%d," % (i), end=None)
    
    # Record the ending time.
    end = time.time()

    print(end - start)

    # Save the image.
    cv2.imwrite("ValidPixels.png", img)

上面的脚本在我的机器上执行时间约520s，图像像素约为1200万。之前早有耳闻说是像我这样简单的进行for循环，python的效率会非常低，今天算是见识到了。于是想着，得，还是改成C++的吧。但是转念一想，貌似有一个叫什么“牛B”的库可以用，可以使用CUDA进行加速。简单查了一下果然是有这个库，numba。

可能是我太闲，我决定试一下。

numba + CUDA

numba可以在没有CUDA支持时使用CPU进行加速，而这里我只感兴趣CUDA的部分。

numba要用conda配置，也是醉了。还好用了conda environment。

我想说numba的文档风格我有点不适应，也许是我看的太粗略，一时间没有参透其中的道理。官方的reference manual之外，还有一个github上的项目对numba的CUDA部分进行了解释。这个项目是gtc2018-numba。

于是我趟了几个坑，这里总结一下。

numba天生支持NumPy，但是CUDA部分仅提供非常有限的支持

在使用numba时，主要数据类型与NumPy非常接近，大部分NumPy的属性和接口都有对应的实现，有些接口的参数数量和种类可能受到限制，这需要仔细看numba的官方文档关于NumPy的一部分。但是中的一个坑是，numba的CUDA部分对python的支持更受限制，这需要参考官方文档的CUDA部分。很容易进的坑包括

• kernel函数内部对大部分NumPy的函数不支持。
  CUDA kernel函数内不允许任何可能产生memory allocation的NumPy函数，这使得绝大部分NumPy函数都不能使用了，尤其是创建array的函数，例如zeros(), array()，和linspace()等。对array进行运算的函数也不支持。但是我们仍然可以使用NumPy的标量函数，例如sin()。kernel函数内部可以使用NumPy对象的shape属性。NumPy的indexing和slicing操作似乎仍然可以使用。
• kernel函数内部对list貌似不支持，对tuple支持。貌似可以创建list，但是索引和更新时会失败。一个在stackoverflow上的讨论。dictionary和generator就不要想了。
• kernel函数内部貌似不支持对自定义类的使用。
• device函数需要使用@cuda.jit(device=True)进行修饰，可以有返回值。我不知道为什么官方和上述github教程对device函数没有专门的章节，也许是我没有找到。
• 更新20200427：numba0.46: 不支持类型转换时复用变量名。例如idx = float(value)，之后idx = int(idx)不会导致jit编译错误，但是实际上idx仍为float类型，无法转换。更换为未使用过的变量名，例如i = int(idx)是可以的，此时i为int类型。

于是过了一些远大于520s的时间以后，如下的测试代码终于可以运行了。

CUDA部分代码

from __future__ import print_function

import cv2
import math
import matplotlib.pyplot as plt
import numpy as np
from numba import cuda
import time

@cuda.jit(device=True)
def d_radius_validate(cx, cy, R, width, x, y):
    x = x - cx
    y = y - cy

    r = math.sqrt( x * x + y * y )

    if ( r >= R - width and r <= R + width ):
        return 255
    else:
        return 0

@cuda.jit(device=True)
def d_polar_line_segment_validate(cx, cy, theta, length, width, x, y):
    n0 = length * math.cos(theta) / length
    n1 = length * math.sin(theta) / length

    v0 = x - cx
    v1 = y - cy

    proj = n0 * v0 + n1 * v1

    if ( proj < 0 ):
        return 0
    elif ( proj > length ):
        return 0
    
    oth0 = v0 - proj*n0
    oth1 = v1 - proj*n1

    d = math.sqrt( oth0 * oth0 + oth1 * oth1 )

    if ( d <= width ):
        return 255
    else:
        return 0

@cuda.jit
def k_validate(imgOut):
    tx = cuda.blockIdx.x*cuda.blockDim.x + cuda.threadIdx.x
    ty = cuda.blockIdx.y*cuda.blockDim.y + cuda.threadIdx.y

    xStride = cuda.blockDim.x * cuda.gridDim.x
    yStride = cuda.blockDim.y * cuda.gridDim.y

    cx, cy = 2056, 1504
    R = 1504*0.75
    halfWidth = 10.0

    for y in range( ty, imgOut.shape[0], yStride ):
        for x in range( tx, imgOut.shape[1], xStride ):
            flag = d_radius_validate( cx, cy, R, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag

            flag = d_polar_line_segment_validate( cx, cy, -math.pi/2, R, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag

            flag = d_polar_line_segment_validate( cx, cy, -math.pi/4, R, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag
            
            flag = d_polar_line_segment_validate( cx-halfWidth, cy, 0.0, R + halfWidth, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag
            
            flag = d_polar_line_segment_validate( cx, cy, math.pi/4, R, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag
            
            flag = d_polar_line_segment_validate( cx, cy, math.pi/2, R, halfWidth, 1.0*x, 1.0*y )
            if ( 0 != flag ):
                imgOut[y, x] = flag

if __name__ == "__main__":
    # Create an image.
    img = np.zeros((3008, 4112), dtype=np.int32)

    # Record the starting time.
    start = time.time()

    dImg = cuda.to_device(img)

    cuda.synchronize()
    k_validate[[100, 100, 1], [16, 16, 1]](dImg)
    cuda.synchronize()

    img = dImg.copy_to_host()

    # Record the ending time.
    end = time.time()

    print(end - start)

    # Save the image.
    cv2.imwrite("ValidPixels_numba.png", img)

    print("Done.")

如上述代码所示，采用100x100的grid，和16x16的block thread配置，在当前的GPU（GeForce GTX 980Ti，拿不上台面啊）上的执行时间约为0.4s。

嗯，快了。

也许我真的是闲的。