技術背景

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更高的效能的。