超过Numpy的速度有多难？试试Numba的GPU加速

DechinPhy

发布于 2021-08-31 17:05:06

1.8K0

发布于 2021-08-31 17:05:06

技术背景

Numpy是在Python中非常常用的一个库，不仅具有良好的接口文档和生态，还具备了最顶级的性能，这个库很大程度上的弥补了Python本身性能上的缺陷。虽然我们也可以自己使用Cython或者是在Python中调用C++的动态链接库，但是我们自己实现的方法不一定有Numpy实现的快，这得益于Numpy对于SIMD等技术的深入实现，把CPU的性能发挥到了极致。因此我们只能考虑弯道超车，尝试下能否用自己实现的GPU的算法来打败Numpy的实现。

矩阵的元素乘

为了便于测试，我们这里使用矩阵的元素乘作为测试的案例。所谓的矩阵元素乘，就是矩阵每一个位置的元素对应相乘，注意区分于矩阵乘法，而我们这里为了节省内存，使用的是计算自身的平方这个案例。

# cuda_test.py

import numpy as np
import time
from numba import cuda
cuda.select_device(1)

@cuda.jit
def CudaSquare(x):
    i, j = cuda.grid(2)
    x[i][j] *= x[i][j]

if __name__ == '__main__':
    np.random.seed(1)
    array_length = 2**10
    random_array = np.random.rand(array_length, array_length)
    random_array_cuda = cuda.to_device(random_array)
    square_array = np.square(random_array)
    CudaSquare[(array_length,array_length),(1,1)](random_array_cuda)
    square_array_cuda = random_array_cuda.copy_to_host()
    print (np.sum(square_array-square_array_cuda))

这个案例主要是通过numba的cuda.jit这一装饰器来实现的GPU加速，在这个装饰器下的函数可以使用CUDA的语法，目前来看应该是最Pythonic的CUDA实现方案，相比于pycuda来说。这个被CUDA装饰的函数，只是将矩阵的每一个元素跟自身相乘，也就是取了一个平方，跟numpy.square的算法实现的是一样的，这里我们可以看看运行结果：

$ python3 cuda_test.py 
0.0

这个打印的结果表示，用numba的cuda方案与用numpy的square函数计算出来的结果差值是0，也就是得到了完全一样的结果。需要注意的是，在GPU上的向量是不能够直接打印出来的，需要先用copy_to_host的方法拷贝到CPU上再进行打印。

numba.cuda加速效果测试

在上一个测试案例中，为了展示结果的一致性，我们使用了内存拷贝的方法，但是实际上我们如果把所有的运算都放在GPU上面来运行的话，就不涉及到内存拷贝，因此这部分的时间在速度测试的过程中可以忽略不计。

# cuda_test.py

import numpy as np
import time
from tqdm import trange
from numba import cuda
cuda.select_device(1)

@cuda.jit
def CudaSquare(x):
    i, j = cuda.grid(2)
    x[i][j] *= x[i][j]

if __name__ == '__main__':
    numpy_time = 0
    numba_time = 0
    test_length = 1000
    for i in trange(test_length):
        np.random.seed(i)
        array_length = 2**10
        random_array = np.random.rand(array_length, array_length)
        random_array_cuda = cuda.to_device(random_array)
        time0 = time.time()
        square_array = np.square(random_array)
        time1 = time.time()
        CudaSquare[(array_length,array_length),(1,1)](random_array_cuda)
        time2 = time.time()
        numpy_time += time1-time0
        numba_time += time2-time1
    print ('The time cost of numpy is {}s for {} loops'.format(numpy_time, test_length))
    print ('The time cost of numba is {}s for {} loops'.format(numba_time, test_length))

在这个案例中，我们循环测试1000次的运行效果，测试对象是1024*1024大小的随机矩阵的平方算法。之所以需要这么多次数的测试，是因为numba的即时编译在第一次执行时会消耗一定的编译时间，但是编译完成后再调用，时间就会被大大的缩减。

$ python3 cuda_test.py 
100%|██████████████████████████████████████| 1000/1000 [00:13<00:00, 76.83it/s]
The time cost of numpy is 1.4523804187774658s for 1000 loops
The time cost of numba is 0.46444034576416016s for 1000 loops

可以看到这个运行效果，我们自己的numba实现相比numpy的实现方案要快上2倍左右。但是我们需要有一个这样的概念，就是对于GPU来说，在显存允许的范围内，运算的矩阵维度越大，加速效果就越明显，因此我们再测试一个更大的矩阵：

# cuda_test.py

import numpy as np
import time
from tqdm import trange
from numba import cuda
cuda.select_device(1)

@cuda.jit
def CudaSquare(x):
    i, j = cuda.grid(2)
    x[i][j] *= x[i][j]

if __name__ == '__main__':
    numpy_time = 0
    numba_time = 0
    test_length = 1000
    for i in trange(test_length):
        np.random.seed(i)
        array_length = 2**12
        random_array = np.random.rand(array_length, array_length)
        random_array_cuda = cuda.to_device(random_array)
        time0 = time.time()
        square_array = np.square(random_array)
        time1 = time.time()
        CudaSquare[(array_length,array_length),(1,1)](random_array_cuda)
        time2 = time.time()
        numpy_time += time1-time0
        numba_time += time2-time1
    print ('The time cost of numpy is {}s for {} loops'.format(numpy_time, test_length))
    print ('The time cost of numba is {}s for {} loops'.format(numba_time, test_length))

这里我们测试了一个4096*4096大小的矩阵的平方算法，可以看到最终的效果如下：

$ python3 cuda_test.py 
100%|████████████████████████████████████████| 100/100 [00:22<00:00,  4.40it/s]
The time cost of numpy is 4.878739595413208s for 100 loops
The time cost of numba is 0.3255774974822998s for 100 loops

在100次的测试中，numba的实现比numpy的实现快了将近15倍！！！

最后，我们可以一起看下中间过程中显卡的使用情况：

因为本机上有2张显卡，日常使用第2张来跑计算任务，因此在代码中设置了cuda.select_device(1)，也就是选择第2块显卡的意思。对于单显卡的用户，这个值应该设置为0.

总结概要

Numpy这个库在Python编程中非常的常用，不仅在性能上补足了Python语言的一些固有缺陷，还具有无与伦比的强大生态。但是即使都是使用Python，Numpy也未必就达到了性能的巅峰，对于我们自己日常中使用到的一些计算的场景，针对性的使用CUDA的功能来进行GPU的优化，是可以达到比Numpy更高的性能的。