Python?提速器numba

更新時間：2022年01月12日 08:28:13 作者：愛摸魚的菜鳥碼農(nóng)

這篇文章主要介紹了Python?提速器numba,相信大部分人都感嘆過python 真的太好用了，但是它真的好慢啊，然而今天我們就來用numba解決Python?慢的這個問題,需要的朋友可以參考一下

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

不過在介紹 numba 之前，我們還是得來看看 python 為什么這么慢：

1.為什么 python 這么慢

用過 python 的人都知道，尤其是在有循環(huán)的情況下，python 會比 C++ 慢很多，所以很多人都避免在 python 代碼里引入復(fù)雜的 for 循環(huán)。我們可以想想 python 和 C++ 寫起來有哪些區(qū)別呢：

動態(tài)變量

如果你寫過 C/C++ 就會發(fā)現(xiàn)，我們需要對變量類型有嚴(yán)格的定義，我們需要定義變量的類型是 int 或者 float 之類的。但是 python 就不一樣了，寫過的 python 的人都知道，它去掉了變量申明和數(shù)據(jù)類型。也就是說，無論啥數(shù)據(jù)，咱啥都不用管，想存就存！那么 python 是如何做到這樣灑脫自由的呢？這就不得不提 python 中萬物皆是對象了，真正的數(shù)據(jù)是存在對象里面的。對于一個簡單的兩個變量的加法，python 每次在做運(yùn)算的時候都得先判斷變量的類型，再取出來進(jìn)行運(yùn)算，而對于 C 來說，簡單的內(nèi)存讀寫和機(jī)器指令 ADD 即可。其實(shí)在 C/C++ 中也有可變數(shù)據(jù)類型，但是其聲明是非常復(fù)雜的，是一種非常令人頭疼的結(jié)構(gòu)。

解釋性語言

C/C++ 這類編譯性語言最大的好處就是其編譯過程是發(fā)生在運(yùn)行之前的，源代碼在調(diào)用前被編譯器轉(zhuǎn)換為可執(zhí)行機(jī)器碼，這樣就節(jié)約了大量的時間。而 python 作為一種解釋性語言，沒法做到一次編譯，后續(xù)可以直接運(yùn)行，每次運(yùn)行的時候都要重新將源代碼通過解釋器轉(zhuǎn)化為機(jī)器碼。這樣一個好處就是非常容易 debug（這里要再次感嘆一下 python 真不愧是新手友好型語言~），當(dāng)然，這個問題自然也是有嘗試解決的辦法，一個很重要的技術(shù)就是 JIT （Just-in-time compilation）：JIT 即時編譯技術(shù)是在運(yùn)行時（runtime）將調(diào)用的函數(shù)或程序段編譯成機(jī)器碼載入內(nèi)存，以加快程序的執(zhí)行。說白了，就是在第一遍執(zhí)行一段代碼前，先執(zhí)行編譯動作，然后執(zhí)行編譯后的代碼。

上面只是簡單列出了兩點(diǎn)，當(dāng)然還有更多的原因，限于篇幅就不再具體介紹，而我們開篇提到的 numba 就是通過 JIT 加速了 python 代碼。那么怎么使用 numba 加速我們的代碼呢？我們可以看一些簡單的例子：

2.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?

當(dāng)需要計算的矩陣很小的時候，貌似速度也不慢，可以接受，但是如果輸入的矩陣大小為 (500, 500)，

a = np.random.random((500, 500))?
%timeit cal_sum(a)?

輸出結(jié)果為：

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)?

輸出結(jié)果為：

236 µs ± 545 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

注意在這里我們使用了%itemit 測試運(yùn)行時間（原因我們留到后面說），通過對比兩個時間，我們可以發(fā)現(xiàn)通過 numba 獲得了非常明顯的加速效果！

我們來具體看一下如何用 numba 加速 python 代碼：在實(shí)際使用過程中，numba 其實(shí)是以裝飾器的形式加在 python 函數(shù)上的，用戶可以不用關(guān)心到底 numba 是通過什么方法來優(yōu)化代碼，只需要調(diào)用就行。同時需要注意到 @jit 裝飾器同時也有一個參數(shù) nopython，這個參數(shù)主要是來區(qū)分 numba 的運(yùn)行模式，numba 其實(shí)有兩種運(yùn)行模式：一個是 nopython 模式，另一個就是 object模式。只有在nopython 模式下，才會獲得最好的加速效果，如果 numba 發(fā)現(xiàn)你的代碼里有它不能理解的東西，就會自動進(jìn)入 object 模式，保證程序至少是能夠運(yùn)行的（當(dāng)然這其實(shí)就失去了添加 numba 的意義）。如果我們將裝飾器改為 @jit(nopython=True) 或者 @njit，numba 會假設(shè)你已經(jīng)對所加速的函數(shù)非常了解，強(qiáng)制使用加速的方式，不會進(jìn)入 object 模式，如編譯不成功，則直接拋出異常。

當(dāng)然說到這里，可能大家還是很困惑，numba 到底是怎么加速 python 代碼的？

python 代碼的編譯過程包括四個階段：詞法分析 -> 語法分析 -> 生成字節(jié)碼 -> 將字節(jié)碼解釋為機(jī)器碼執(zhí)行, 常見的 python 解釋器的類型有 cpython、IPython、PyPy、Jython、IronPython，與其他解釋器不同，numba 是使用 LLVM 編譯技術(shù)來解釋字節(jié)碼的。

LLVM 是一個編譯器，它采用字節(jié)碼，并將其編譯為機(jī)器碼，編譯過程涉及許多額外的傳遞，而 LLVM編譯器可以優(yōu)化字節(jié)碼，例如某些頻繁執(zhí)行的模塊，LLVM 可以將其作為 “hot code” 從而進(jìn)行相應(yīng)的優(yōu)化，LLVM 工具鏈非常擅長優(yōu)化字節(jié)碼，它不僅可以編譯 numba 的代碼，還可以優(yōu)化它。

在第一次調(diào)用 numba 裝飾的函數(shù)時，numba 將在調(diào)用期間推斷參數(shù)類型，numba 會結(jié)合給定的參數(shù)類型將其編譯為機(jī)器代碼。這個過程是有一定的時間消耗的，但是一旦編譯完成，numba 會為所呈現(xiàn)的特定類型的參數(shù)緩存函數(shù)的機(jī)器代碼版本，如果再次使用相同的類型調(diào)用它，它可以重用緩存的機(jī)器代碼而不必再次編譯。

在測量性能時，如果只使用一個簡單的計時器來計算一次，該計時器包括在執(zhí)行時編譯函數(shù)所花費(fèi)的時間，最準(zhǔn)確的運(yùn)行時間應(yīng)該是第二次及以后調(diào)用函數(shù)的運(yùn)行時間。
對于指定輸入類型這個問題，我們可以嘗試做一個簡單的實(shí)驗(yàn)看看到底有怎樣的影響：

a = np.random.random((5000, 5000))?
?
# 第一次調(diào)用時間包括編譯時間?
start = time.time()?
cal_sum(a)?
end = time.time()?
print("Elapsed (with compilation) = %s" % (end - start))?
?
# 函數(shù)被編譯，機(jī)器代碼被緩存?
start = time.time()?
cal_sum(a)?
end = time.time()?
print("Elapsed (after compilation) = %s" % (end - start))?
?
# 這里 a 本身的類型為 np.float64?
b = a.astype(np.float32)?
?
# 調(diào)用相同的函數(shù)，但是輸入數(shù)據(jù)的類型變?yōu)?np.float32?
start = time.time()?
cal_sum(b)?
end = time.time()?
print("Elapsed (after compilation) = %s" % (end - start))?

輸出結(jié)果：

Elapsed (with compilation) = 0.20406198501586914
Elapsed (after compilation) = 0.025263309478759766
Elapsed (after compilation) = 0.07892274856567383

可以看到如果我們輸入了和第一次調(diào)用編譯時不同的數(shù)據(jù)類型，函數(shù)的運(yùn)行時間也會有一個很明顯的增加，但仍然是遠(yuǎn)低于第一次運(yùn)行時的編譯的時間。

3. 如果調(diào)用 numba 的時候顯式地指定輸入、輸出數(shù)據(jù)的類型，可以加快初次調(diào)用的函數(shù)時的編譯速度，同時壞處就是如果顯式指定后，那么之后調(diào)用該函數(shù)都必須滿足規(guī)定的數(shù)據(jù)類型。

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?
?
# 不指定輸入輸出數(shù)據(jù)類型，讓 numba 自己判斷?
start = time.time()?
cal_sum1(a)?
end = time.time()?
print("Elapsed (with compilation) = %s" % (end - start))?
?
# 指定輸入輸出數(shù)據(jù)類型?
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?

可以看到編譯的時間被大大減少了，其實(shí)這個時間非常接近直接運(yùn)行該函數(shù)生成的機(jī)器代碼的時間。

上面說了這么多，但是轉(zhuǎn)念一想，矩陣相加這個函數(shù) numpy 里好像早就有了，np.sum 它不好用，它不香嘛？？干嘛搞得這么復(fù)雜？

好吧，就上面舉的簡單的例子來說，使用 numpy 和 numba 加速基本效果差不多，但是在實(shí)際情況里面，不是所有的 for 循環(huán)代碼都可以直接用 numpy 自帶的函數(shù)實(shí)現(xiàn)。但是 numba 基本對所有的 for 循環(huán)代碼都有非常好的加速效果，當(dāng)然前提是 for 循環(huán)里面的代碼必須是 numba 能夠理解的。

而在從實(shí)際使用中，一般推薦將代碼中密集的計算部分提取出來作為單獨(dú)的函數(shù)實(shí)現(xiàn)，并使用 nopython 方式優(yōu)化，這樣可以保證我們能使用到 numba 的加速功能。其余部分還是使用 python 原生代碼，這樣一方面就可以做到在 numba 加速不明顯或者無法加速的代碼中調(diào)用各種函數(shù)實(shí)現(xiàn)自己的代碼邏輯，另一方面也能享受到 numba 的加速效果。

3.numba 加速 numpy 運(yùn)算

上面說了 numba 一大亮點(diǎn)就是加速 for 循環(huán)，除此以外，numba 對 numpy 的運(yùn)算也同樣的有加速的效果。因?yàn)榧词故?numpy 也沒有 numba 轉(zhuǎn)換為機(jī)器碼快，numba 尤其擅長加速 numpy 的基本運(yùn)算 (如加法、相乘和平方等等) ，其實(shí)準(zhǔn)確來說如果 numpy 函數(shù)是對各個元素采用相同的操作的情況下，都會有比較好的效果。

我們簡單舉一個 numba 加速 numpy 運(yùn)算的例子：

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)?
?
# 第一次調(diào)用完成編譯?
add_arrays_numba(a)?
?
# 函數(shù)被編譯，機(jī)器代碼被緩存?
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

當(dāng)我們對 numpy 數(shù)組進(jìn)行基本的數(shù)組計算，比如加法、乘法和平方，numpy 都會自動在內(nèi)部向量化，這也是它可以比原生 python 代碼有更好性能的原因。但是在特定情況下，numpy 的代碼也不會和優(yōu)化過的機(jī)器代碼速度一樣快，此時 numba 直接作用于 numpy 運(yùn)算也能起到一定的加速效果。

另一個例子主要來自于MMDetection3D，經(jīng)過一定的簡化，主要是用來計算將點(diǎn)的坐標(biāo) (x, y) 壓縮到給定的[x_min, y_min, x_max, y_max] 范圍內(nèi):

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)?

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

4.numba 使用 CUDA 加速

numba 更厲害的地方就在于，我們可以直接用 python 寫 CUDA Kernel，直接在 GPU 上編譯和運(yùn)行我們的 Python 程序，numba 通過將 python 代碼直接編譯為遵循 CUDA 執(zhí)行模型的 CUDA 內(nèi)核和設(shè)備函數(shù)來支持 CUDA GPU 編程（但是實(shí)際上 numba 目前支持的 CUDA API 很少，希望開發(fā)團(tuán)隊能更肝一點(diǎn)~~~) ，為了節(jié)省將 numpy 數(shù)組復(fù)制到指定設(shè)備，然后又將結(jié)果存儲到 numpy 數(shù)組中所浪費(fèi)的時間，numba 提供了一些函數(shù)來聲明并將數(shù)組送到指定設(shè)備來節(jié)省不必要的復(fù)制到 cpu 的時間。

常用內(nèi)存分配函數(shù)：

cuda.device_array()：在設(shè)備上分配一個空向量，類似于numpy.empty()；
cuda.to_device()：將主機(jī)的數(shù)據(jù)拷貝到設(shè)備；
cuda.copy_to_host()：將設(shè)備的數(shù)據(jù)拷貝回主機(jī)；

我們可以通過一個簡單的矩陣相加的例子來看看通過 numba 使用 CUDA 加速的效果：

from numba import cuda # 從numba調(diào)用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)
?
# 拷貝數(shù)據(jù)到設(shè)備端
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)
?
# 在顯卡設(shè)備上初始化一塊用于存放GPU計算結(jié)果的空間
gpu_result1 = cuda.device_array((m,n))
gpu_result2 = 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)
# 第一次調(diào)用包含編譯時間
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result1, m, n)
cuda.synchronize()
print("gpu matrix add time (with compilation) " + str(time() - start))
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result2, m, n)
cuda.synchronize()
print("gpu matrix add time (after compilation)" + str(time() - start))
start = time()
cpu_result = np.add(x, y)
print("cpu matrix add time " + str(time() - start))

運(yùn)行時間分別為：

gpu matrix add time (with compilation) 0.15977692604064941
gpu matrix add time (after compilation) 0.0005376338958740234
cpu matrix add time 0.023023128509521484

在通過 numba 進(jìn)行 CUDA 加速的時候，主要是通過調(diào)用@cuda.jit 裝飾器實(shí)現(xiàn)，從結(jié)果可以看到 numba 通過調(diào)用 CUDA 明顯加速了 python 程序。

5.For 循環(huán)寫法的影響

下面的一段代碼截取自MMDetection3D, 主要是用來判斷一系列點(diǎn)是否在一系列多邊形的內(nèi)部，

我們可以有如下的兩種寫法：

在 For 循環(huán)里面計算 vec1，每次循環(huán)都需要訪問多邊形 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?

在循環(huán)前預(yù)先計算好所有的 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

通過測試我們可以發(fā)現(xiàn)第二種方案會更快，在實(shí)際使用的時候，我們可以盡量減少在 for 循環(huán)內(nèi)部內(nèi)存的訪問次數(shù)，從而降低函數(shù)的運(yùn)行時間。

總結(jié) ：

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

到此這篇關(guān)于Python 提速器numba的文章就介紹到這了,更多相關(guān)Python numba內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: