技术背景
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更高的性能的。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/numba-cuda.html
作者ID：DechinPhy
更多原著文章请参考：https://www.cnblogs.com/dechinphy/
打赏专用链接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html