技術背景
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