Python 提速大殺器之 numba 篇教程

你是不是曾經有這樣的苦惱，python 真的太好用了，但是它真的好慢?。匏? ; C++ 很快，但是真的好難寫啊，此生能不碰它就不碰它。老天啊，有沒有什么兩全其美的辦法呢？俗話說的好：辦法總是比困難多，大家都有這個問題，自然也就有大佬來試著解決這個問題，這就請出我們今天的主角: numba

不過在介紹 numba 之前，我們還是得來看看 python 為什么這么慢：
為什么 python 這么慢
用過 python 的人都知道， 尤其是在有循環的情況下，python 會比 C++ 慢很多，所以很多人都避免在 python 代碼里引入復雜的 for 循環。我們可以想想 python 和 C++ 寫起來有哪些區別呢：
動態變量
如果你寫過 C/C++ 就會發現，我們需要對變量類型有嚴格的定義，我們需要定義變量的類型是 int 或者 float 之類的。但是 python 就不一樣了，寫過的 python 的人都知道，它去掉了變量申明和數據類型。也就是說，無論啥數據，咱啥都不用管，想存就存！那么 python 是如何做到這樣灑脫自由的呢？這就不得不提 python 中萬物皆是對象了，真正的數據是存在對象里面的。對于一個簡單的兩個變量的加法，python 每次在做運算的時候都得先判斷變量的類型，再取出來進行運算，而對于 C 來說，簡單的內存讀寫和機器指令 ADD 即可。其實在 C/C++ 中也有可變數據類型，但是其聲明是非常復雜的，是一種非常令人頭疼的結構。
解釋性語言
C/C++ 這類編譯性語言最大的好處就是其編譯過程是發生在運行之前的，源代碼在調用前被編譯器轉換為可執行機器碼，這樣就節約了大量的時間。而 python 作為一種解釋性語言，沒法做到一次編譯，后續可以直接運行，每次運行的時候都要重新將源代碼通過解釋器轉化為機器碼。這樣一個好處就是非常容易 debug（ 這里要再次感嘆一下 python 真不愧是新手友好型語言~）， 當然，這個問題自然也是有嘗試解決的辦法，一個很重要的技術就是 JIT （Just-in-time compilation）：JIT 即時編譯技術是在運行時（runtime）將調用的函數或程序段編譯成機器碼載入內存，以加快程序的執行。說白了，就是在第一遍執行一段代碼前，先執行編譯動作，然后執行編譯后的代碼。
上面只是簡單列出了兩點，當然還有更多的原因，限于篇幅就不再具體介紹，而我們開篇提到的 numba 就是通過 JIT 加速了 python 代碼。那么怎么使用 numba 加速我們的代碼呢？我們可以看一些簡單的例子：
numba 加速 python 的小例子
用 numba 加速 python 代碼多簡單方便呢，我們先來看看如何使用 numba 加速 python 代碼。
如果讓你用單純的 python 計算一個矩陣所有元素的和，很容易可以寫出下面的代碼：
def cal_sum(a):
    result = 0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            result += a[i, j]
    return result
當需要計算的矩陣很小的時候，貌似速度也不慢，可以接受，但是如果輸入的矩陣大小為 (500, 500)，
a = np.random.random((500, 500))
%timeit cal_sum(a)
輸出結果為：
47.8 ms ± 499 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)
我們嘗試加上 numba：
import numba 

@numba.jit(nopython=True)
def cal_sum(a):
    result = 0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            result += a[i, j]

    return result

輸入同樣大小的矩陣

a = np.random.random((500, 500))
%timeit cal_sum(a)
輸出結果為：
236 μs ± 545 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
注意！在這里我們使用了 % itemit 測試運行時間（原因我們留到后面說），通過對比兩個時間，我們可以發現通過 numba 獲得了非常明顯的加速效果！
我們來具體看一下如何用 numba 加速 python 代碼：在實際使用過程中，numba 其實是以裝飾器的形式加在 python 函數上的，用戶可以不用關心到底 numba 是通過什么方法來優化代碼，只需要調用就行。同時需要注意到 @jit 裝飾器同時也有一個參數 nopython， 這個參數主要是來區分 numba 的運行模式，numba 其實有兩種運行模式：一個是 nopython 模式，另一個就是 object 模式。只有在 nopython 模式下，才會獲得最好的加速效果，如果 numba 發現你的代碼里有它不能理解的東西，就會自動進入 object 模式，保證程序至少是能夠運行的（當然這其實就失去了添加 numba 的意義）。如果我們將裝飾器改為 @jit(nopython=True) 或者 @njit，numba 會假設你已經對所加速的函數非常了解，強制使用加速的方式，不會進入 object 模式，如編譯不成功，則直接拋出異常。
當然說到這里，可能大家還是很困惑，numba 到底是怎么加速 python 代碼的？
python 代碼的編譯過程包括四個階段：詞法分析 -> 語法分析 -> 生成字節碼 -> 將字節碼解釋為機器碼執行, 常見的 python 解釋器的類型有 cpython、IPython、PyPy、Jython、IronPython，與其他解釋器不同，numba 是使用 LLVM 編譯技術來解釋字節碼的。
LLVM 是一個編譯器，它采用字節碼，并將其編譯為機器碼，編譯過程涉及許多額外的傳遞，而 LLVM編譯器可以優化字節碼，例如某些頻繁執行的模塊，LLVM 可以將其作為 “hot code” 從而進行相應的優化，LLVM 工具鏈非常擅長優化字節碼，它不僅可以編譯 numba 的代碼，還可以優化它。
在第一次調用 numba 裝飾的函數時，numba 將在調用期間推斷參數類型，numba 會結合給定的參數類型將其編譯為機器代碼。這個過程是有一定的時間消耗的，但是一旦編譯完成，numba 會為所呈現的特定類型的參數緩存函數的機器代碼版本，如果再次使用相同的類型調用它，它可以重用緩存的機器代碼而不必再次編譯。
- 在測量性能時，如果只使用一個簡單的計時器來計算一次，該計時器包括在執行時編譯函數所花費的時間，最準確的運行時間應該是第二次及以后調用函數的運行時間。
- 對于指定輸入類型這個問題，我們可以嘗試做一個簡單的實驗看看到底有怎樣的影響：
a = np.random.random((5000, 5000))
# 第一次調用時間包括編譯時間
start = time.time()
cal_sum(a)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))
 
# 函數被編譯，機器代碼被緩存
start = time.time()
cal_sum(a)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

# 這里 a 本身的類型為 np.float64
b = a.astype(np.float32)

# 調用相同的函數，但是輸入數據的類型變為 np.float32
start = time.time()
cal_sum(b)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))
輸出結果：
Elapsed (with compilation) = 0.20406198501586914
Elapsed (after compilation) = 0.025263309478759766
Elapsed (after compilation) = 0.07892274856567383
可以看到如果我們輸入了和第一次調用編譯時不同的數據類型，函數的運行時間也會有一個很明顯的增加，但仍然是遠低于第一次運行時的編譯的時間。
- 如果調用 numba 的時候顯式地指定輸入、輸出數據的類型，可以加快初次調用的函數時的編譯速度，同時壞處就是如果顯式指定后，那么之后調用該函數都必須滿足規定的數據類型。
a = np.random.random((500, 500))
@numba.njit()
def cal_sum1(a):
    result = 0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            result += a[i, j]
    return result

@numba.njit('float64(float64[:, :])')
def cal_sum2(a):
    result = 0
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            result += a[i, j]
    return result

# 不指定輸入輸出數據類型，讓 numba 自己判斷
start = time.time()
cal_sum1(a)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

# 指定輸入輸出數據類型
start = time.time()
cal_sum2(a)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

分別耗時：
Elapsed (after compilation) = 0.054465532302856445
Elapsed (after compilation) = 0.0004112720489501953
可以看到編譯的時間被大大減少了，其實這個時間非常接近直接運行該函數生成的機器代碼的時間。
上面說了這么多，但是轉念一想，矩陣相加這個函數 numpy 里好像早就有了，np.sum 它不好用，它不香嘛？？干嘛搞得這么復雜？
好吧，就上面舉的簡單的例子來說，使用 numpy 和 numba 加速基本效果差不多，但是在實際情況里面，不是所有的 for 循環代碼都可以直接用 numpy 自帶的函數實現。但是 numba 基本對所有的 for 循環代碼都有非常好的加速效果，當然前提是 for 循環里面的代碼必須是 numba 能夠理解的。
而在從實際使用中，一般推薦將代碼中密集的計算部分提取出來作為單獨的函數實現，并使用 nopython 方式優化，這樣可以保證我們能使用到 numba 的加速功能。其余部分還是使用 python 原生代碼，這樣一方面就可以做到在 numba 加速不明顯或者無法加速的代碼中調用各種函數實現自己的代碼邏輯， 另一方面也能享受到 numba 的加速效果。
numba 加速 numpy 運算
上面說了 numba 一大亮點就是加速 for 循環，除此以外，numba 對 numpy 的運算也同樣的有加速的效果。因為即使是 numpy 也沒有 numba 轉換為機器碼快，numba 尤其擅長加速 numpy 的基本運算 (如加法、相乘和平方等等) ，其實準確來說如果 numpy 函數是對各個元素采用相同的操作的情況下，都會有比較好的效果。我們簡單舉一個 numba 加速 numpy 運算的例子：
a = np.ones((1000, 1000), np.int64) * 5
b = np.ones((1000, 1000), np.int64) * 10
c = np.ones((1000, 1000), np.int64) * 15

def add_arrays(a, b, c):
    return np.square(a, b, c)

@numba.njit
def add_arrays_numba(a, b, c):
    return np.square(a, b, c)
    
# 第一次調用完成編譯
add_arrays_numba(a)
 
# 函數被編譯，機器代碼被緩存
start = time.time()
add_arrays_numba(a)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

# 不使用 numba 加速
start = time.time()
add_arrays(a)
end = time.time()
print("Elapsed = %s" % (end - start))

Elapsed (after compilation) = 0.002088785171508789
Elapsed = 0.0031290054321289062

當我們對 numpy 數組進行基本的數組計算，比如加法、乘法和平方，numpy 都會自動在內部向量化，這也是它可以比原生 python 代碼有更好性能的原因。但是在特定情況下，numpy 的代碼也不會和優化過的機器代碼速度一樣快，此時 numba 直接作用于 numpy 運算也能起到一定的加速效果。
另一個例子主要來自于 MMDetection3D，經過一定的簡化，主要是用來計算將點的坐標 (x, y) 壓縮到給定的 [x_min, y_min, x_max, y_max] 范圍內:
x = np.random.random((5000))*5000
y = np.random.random((5000))*5000
x_min = 0
x_max = 1000
y_min=0
y_max=2000

@numba.njit
def get_clip_numba(x, y, x_min, y_min, x_max, y_max):
    z = np.stack((x, y), axis=1)
    z[:, 0] = np.clip(z[:, 0], x_min, x_max)
    z[:, 1] = np.clip(z[:, 1], y_min, y_max)
    return z

def get_clip(x, y, x_min, y_min, x_max, y_max):
    z = np.stack((x, y), axis=1)
    z[:, 0] = np.clip(z[:, 0], x_min, x_max)
    z[:, 1] = np.clip(z[:, 1], y_min, y_max)
    return z

%timeit get_clip_numba(x, y, x_min, y_min, x_max, y_max)
%timeit get_clip(x, y, x_min, y_min, x_max, y_max)

分別用時：
33.8 μs ± 12.2 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
57.2 μs ± 258 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

從實際情況來看， 并不是所有的 numpy 函數在使用 numba 后都能獲得比較好的加速效果，在某些情況下甚至會降低 numpy 的運行速度。因此，在實際使用過程中建議提前測試一下確認加速效果。通常將 numba 用于加速 numpy 的時候都是 for 循環和 numpy 一起使用的情況。numba 對 numpy 的大部分常用的函數都做了支持。

numba 使用 CUDA 加速
numba 更厲害的地方就在于，我們可以直接用 python 寫 CUDA Kernel， 直接在 GPU 上編譯和運行我們的 Python 程序，numba 通過將 python 代碼直接編譯為遵循 CUDA 執行模型的 CUDA 內核和設備函數來支持 CUDA GPU 編程（ 但是實際上 numba 目前支持的 CUDA  API 很少，希望開發團隊能更肝一點~~~) ，對于不熟悉 CUDA 的同學， 我們推薦大家看一下之前的一篇文章 PyTorch 源碼解讀之 cpp_extension：揭秘 C++/CUDA 算子實現和調用全流程 （網址：https://zhuanlan.zhihu.com/p/348555597）了解一下 CUDA 編程的基本概念。為了節省將 numpy 數組復制到指定設備，然后又將結果存儲到 numpy 數組中所浪費的時間，numba 提供了一些函數來聲明并將數組送到指定設備來節省不必要的復制到 cpu 的時間。

常用內存分配函數：
- cuda.device_array()：在設備上分配一個空向量，類似于numpy.empty()；
- cuda.to_device()：將主機的數據拷貝到設備；
- cuda.copy_to_host()：將設備的數據拷貝回主機；
我們可以通過一個簡單的矩陣相加的例子來看看通過 numba 使用 CUDA 加速的效果：
from numba import cuda
import numpy as np
import math
from time import time

@cuda.jit
def matrix_add(a, b, result, m, n):
    idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    idy = cuda.threadIdx.y+ cuda.blockDim.y * cuda.blockIdx.y
    if idx < m and idy < n:
        result[idx, idy] = a[idx, idy] + b[idx, idy]
m = 5000
n = 4000

x = np.arange(m*n).reshape((m,n)).astype(np.int32)
y = np.arange(m*n).reshape((m,n)).astype(np.int32)

# 拷貝數據到 gpu 
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)

# 在 gpu 上初始化一塊用于存放 gpu 計算結果的空間
gpu_result = cuda.device_array((m,n))
cpu_result = np.empty((m,n))

threads_per_block = 1024
blocks_per_grid = math.ceil(m*n / threads_per_block)
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result, m, n)
cuda.synchronize()
print("gpu matrix add time " + str(time() - start))
start = time()
cpu_result = np.add(x, y)
print("cpu matrix add time " + str(time() - start))

if (np.array_equal(cpu_result, gpu_result.copy_to_host())):
    print("result correct!")

運行時間分別為：
gpu matrix add time 0.32171130180358887
cpu matrix add time 0.017357587814331055

在通過 numba 進行 CUDA 加速的時候，主要是通過調用 @cuda.jit 裝飾器實現，從結果可以看到 numba 通過調用 CUDA 明顯加速了 python 程序。
For 循環寫法的影響
下面的一段代碼截取自 MMDetection3D ,  主要是用來判斷一系列點是否在一系列多邊形的內部，我們可以有如下的兩種寫法：
- 在 For 循環里面計算 vec1， 每次循環都需要訪問多邊形 polygon 變量
@numba.jit(nopython=True)
def points_in_convex_polygon1(points, polygon, clockwise=True):
    # first convert polygon to directed lines
    num_points_of_polygon = polygon.shape[1]
    num_points = points.shape[0]
    num_polygons = polygon.shape[0]
    vec1 = np.zeros((2), dtype=polygon.dtype)
    ret = np.zeros((num_points, num_polygons), dtype=np.bool_)
    success = True
    cross = 0.0
    for i in range(num_points):
        for j in range(num_polygons):
            success = True
            for k in range(num_points_of_polygon):
                if clockwise:
                    vec1 = polygon[j, k] - polygon[j, k - 1]
                else:
                    vec1 = polygon[j, k - 1] - polygon[j, k]
                cross = vec1[1] * (polygon[j, k, 0] - points[i, 0])
                cross -= vec1[0] * (polygon[j, k, 1] - points[i, 1])
                if cross >= 0:
                    success = False
                    break
            ret[i, j] = success
    return ret


- 在循環前預先計算好所有的 vec
@numba.jit(nopython=True)
def points_in_convex_polygon2(points, polygon, clockwise=True):
    # first convert polygon to directed lines
    num_points_of_polygon = polygon.shape[1]
    num_points = points.shape[0]
    num_polygons = polygon.shape[0]
    # vec for all the polygons
    if clockwise:
        vec1 = polygon - polygon[:, np.array([num_points_of_polygon - 1] +
                                 list(range(num_points_of_polygon - 1))), :]
    else:
        vec1 = polygon[:, np.array([num_points_of_polygon - 1] +
                       list(range(num_points_of_polygon - 1))), :] - polygon
    ret = np.zeros((num_points, num_polygons), dtype=np.bool_)
    success = True
    cross = 0.0
    for i in range(num_points):
        for j in range(num_polygons):
            success = True
            for k in range(num_points_of_polygon):
                vec = vec1[j,k]
                cross = vec[1] * (polygon[j, k, 0] - points[i, 0])
                cross -= vec[0] * (polygon[j, k, 1] - points[i, 1])
                if cross >= 0:
                    success = False
                    break
            ret[i, j] = success
    return ret

簡單測試一下兩種寫法的速度：
points = np.random.random((20000, 2)) * 100
polygon = np.random.random((1000, 100, 2)) * 200 

start = time.time()
points_in_convex_polygon1(points, polygon)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

start = time.time()
points_in_convex_polygon1(points, polygon)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

start = time.time()
points_in_convex_polygon2(points, polygon)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

start = time.time()
points_in_convex_polygon2(points, polygon)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))


輸出時間：
Elapsed (with compilation) = 3.9232356548309326
Elapsed (after compilation) = 3.6778993606567383
Elapsed (with compilation) = 0.6269152164459229
Elapsed (after compilation) = 0.22288227081298828


通過測試我們可以發現第二種方案會更快，在實際使用的時候，我們可以盡量減少在 for 循環內部內存的訪問次數，從而降低函數的運行時間。

總結 + 一點碎碎念

我們介紹了一些用 numba 加速的常見場景，能夠有效地提高我們代碼的速度。不過大家在使用的時候，建議多多嘗試，比較一下使用與不使用的速度區別（有時候用了 numba 還可能變得更慢......)，此外 MMDetection3D 很早就使用了 numba 加速代碼，而且我們最近在 MMDetection3D 中升級了 numba 的版本，從而獲得更好的 numpy 兼容性和代碼加速效果。

以上就是“Python 提速大殺器之 numba 篇教程”的詳細內容，想要了解更多Python教程歡迎持續關注編程學習網。