“brown snake” by Duncan Sanchez on Unsplash
這是我寫的該系列文章中的第二篇。所有的文章如下:
這些文章對應的 Jupyter Notebooks 程式碼如下:
[Github-SpeedUpYourAlgorithms] and [Kaggle]
目錄
- 介紹
- 為什麼選擇 Numba?
- Numba 是如何工作的?
- 使用 Numba 的基本功能(只需要加上
@jit
!) - @vectorize 裝飾器
- 在 GPU 上執行函式
- 擴充套件閱讀
- 參考
注意: 這篇文章的 Jupyter Notebook 程式碼在我的 Github 上:SpeedUpYourAlgorithms-Numba
1. 介紹
Numba 是 python 的即時(Just-in-time)編譯器,即當您呼叫 python 函式時,您的全部或部分程式碼就會被轉換為“即時”執行的機器碼,它將以您的本地機器碼速度執行!它由 Anaconda 公司贊助,並得到了許多其他組織的支援。
在 Numba 的幫助下,您可以加速所有計算負載比較大的 python 函式(例如迴圈)。它還支援 numpy 庫!所以,您也可以在您的計算中使用 numpy,並加快整體計算,因為 python 中的迴圈非常慢。 您還可以使用 python 標準庫中的 math 庫的許多函式,如 sqrt
等。有關所有相容函式的完整列表,請檢視 此處。
2. 為什麼選擇 Numba?
那麼,當有像 cython 和 Pypy 之類的許多其他編譯器時,為什麼要選擇 numba?
原因很簡單,這樣您就不必離開寫 python 程式碼的舒適區。是的,就是這樣,您根本不需要為了獲得一些的加速來改變您的程式碼,這與您從類似的具有型別定義的 cython 程式碼獲得的加速相當。那不是很好嗎?
您只需要新增一個熟悉的 python 功能,即新增一個包裝器(一個裝飾器)到您的函式上。類的裝飾器也在開發中了。
所以,您只需要新增一個裝飾器就可以了。例如:
1 2 3 4 5 6 |
from numba import jit @jit def function(x): # your loop or numerically intensive computations return x |
這仍然看起來像一個原生 python 程式碼,不是嗎?
3. 如何使用 Numba?
“question mark neon signage” by Emily Morter on Unsplash
Numba 使用 LLVM 編譯器基礎結構 將原生 python 程式碼轉換成優化的機器碼。使用 numba 執行程式碼的速度可與 C/C++ 或 Fortran 中的類似程式碼相媲美。
以下是程式碼的編譯方式:
首先,Python 函式被傳入,優化並轉換為 numba 的中間表達,然後在型別推斷(type inference)之後,就像 numpy 的型別推斷(所以 python float 是一個 float64),它被轉換為 LLVM 可解釋程式碼。 然後將此程式碼提供給 LLVM 的即時編譯器以生成機器碼。
您可以根據需要在執行時或匯入時 生成 機器碼,匯入需要在 CPU(預設)或 GPU 上進行。
4. 使用 numba 的基本功能(只需要加上 @jit !)
Photo by Charles Etoroma on Unsplash
小菜一碟!
為了獲得最佳效能,numba 實際上建議在您的 jit 裝飾器中加上 nopython=True
引數,加上後就不會使用 Python 直譯器了。或者您也可以使用 @njit
。如果您加上 nopython=True
的裝飾器失敗並報錯,您可以用簡單的 @jit
裝飾器來編譯您的部分程式碼,對於它能夠編譯的程式碼,將它們轉換為函式,並編譯成機器碼。然後將其餘部分程式碼提供給 python 直譯器。
所以,您只需要這樣做:
1 2 3 4 5 6 |
from numba import njit, jit @njit # or @jit(nopython=True) def function(a, b): # your loop or numerically intensive computations return result |
當使用 @jit
時,請確保您的程式碼有 numba 可以編譯的內容,比如包含庫(numpy)和它支援的函式的計算密集型迴圈。否則它將不會編譯任何東西,並且您的程式碼將比沒有使用 numba 時更慢,因為存在 numba 內部程式碼檢查的額外開銷。
還有更好的一點是,numba 會對首次作為機器碼使用後的函式進行快取。 因此,在第一次使用之後它將更快,因為它不需要再次編譯這些程式碼,如果您使用的是和之前相同的引數型別。
如果您的程式碼是 可並行化 的,您也可以傳遞 parallel=True
作為引數,但它必須與 nopython=True
一起使用,目前這隻適用於CPU。
您還可以指定希望函式具有的函式簽名,但是這樣就不會對您提供的任何其他型別的引數進行編譯。 例如:
1 2 3 4 5 6 7 8 9 10 11 12 |
from numba import jit, int32 @jit(int32(int32, int32)) def function(a, b): # your loop or numerically intensive computations return result # or if you haven't imported type names # you can pass them as string @jit('int32(int32, int32)') def function(a, b): # your loop or numerically intensive computations return result |
現在您的函式只能接收兩個 int32 型別的引數並返回一個 int32 型別的值。 通過這種方式,您可以更好地控制您的函式。 如果需要,您甚至可以傳遞多個函式簽名。
您還可以使用 numba 提供的其他裝飾器:
- @vectorize:允許將標量引數作為 numpy 的 ufuncs 使用,
- @guvectorize:生成 NumPy 廣義上的
ufunc
s, - @stencil:定義一個函式使其成為 stencil 型別操作的核函式
- @jitclass:用於 jit 類,
- @cfunc:宣告一個函式用於本地回撥(被C/C++等呼叫),
- @overload:註冊您自己的函式實現,以便在
nopython
模式下使用,例如:@overload(scipy.special.j0)
。
Numba 還有 Ahead of time(AOT)編譯,它生成不依賴於 Numba 的已編譯擴充套件模組。 但:
- 它只允許常規函式(ufuncs 就不行),
- 您必須指定函式簽名。並且您只能指定一種簽名,如果需要指定多個簽名,需要使用不同的名字。
它還根據您的CPU架構系列生成通用程式碼。
5. @vectorize 裝飾器
“gray solar panel lot” by American Public Power Association on Unsplash
通過使用 @vectorize 裝飾器,您可以對僅能對標量操作的函式進行轉換,例如,如果您使用的是僅適用於標量的 python 的 math
庫,則轉換後就可以用於陣列。 這提供了類似於 numpy 陣列運算(ufuncs)的速度。 例如:
1 2 3 4 5 |
@vectorize def func(a, b): # Some operation on scalars return result |
您還可以將 target
引數傳遞給此裝飾器,該裝飾器使 target 引數為 parallel
時用於並行化程式碼,為 cuda
時用於在 cuda\GPU 上執行程式碼。
1 2 3 4 5 |
@vectorize(target="parallel") def func(a, b): # Some operation on scalars return result |
使 target=“parallel”
或 “cuda”
進行向量化通常比 numpy 實現的程式碼執行得更快,只要您的程式碼具有足夠的計算密度或者陣列足夠大。如果不是,那麼由於建立執行緒以及將元素分配到不同執行緒需要額外的開銷,因此可能耗時更長。所以運算量應該足夠大,才能獲得明顯的加速。
這個視訊講述了一個用 Numba 加速用於計算流體動力學的Navier Stokes方程的例子:
6. 在GPU上執行函式
“time-lapsed of street lights” by Marc Sendra martorell on Unsplash
您也可以像裝飾器一樣傳遞 @jit 來執行 cuda/GPU 上的函式。 為此您必須從 numba
庫中匯入 cuda
。 但是要在 GPU 上執行程式碼並不像之前那麼容易。為了在 GPU 上的數百甚至數千個執行緒上執行函式,需要先做一些初始計算。 實際上,您必須宣告並管理網格,塊和執行緒的層次結構。這並不那麼難。
要在GPU上執行函式,您必須定義一個叫做 核函式 或 裝置函式 的函式。首先讓我們來看 核函式。
關於核函式要記住一些要點:
a)核函式在被呼叫時要顯式宣告其執行緒層次結構,即塊的數量和每塊的執行緒數量。您可以編譯一次核函式,然後用不同的塊和網格大小多次呼叫它。
b)核函式沒有返回值。因此,要麼必須對原始陣列進行更改,要麼傳遞另一個陣列來儲存結果。為了計算標量,您必須傳遞單元素陣列。
1 2 3 4 5 6 7 8 |
# Defining a kernel function from numba import cuda @cuda.jit def func(a, result): # Some cuda related computation, then # your computationally intensive code. # (Your answer is stored in 'result') |
因此,要啟動核函式,您必須傳入兩個引數:
- 每塊的執行緒數,
- 塊的數量。
例如:
1 2 3 4 |
threadsperblock = 32 blockspergrid = (array.size + (threadsperblock - 1)) // threadsperblock func[blockspergrid, threadsperblock](array) |
每個執行緒中的核函式必須知道它在哪個執行緒中,以便了解它負責陣列的哪些元素。Numba 只需呼叫一次即可輕鬆獲得這些元素的位置。
1 2 3 4 5 6 7 |
@cuda.jit def func(a, result): pos = cuda.grid(1) # For 1D array # x, y = cuda.grid(2) # For 2D array if pos < a.shape[0]: result[pos] = a[pos] * (some computation) |
為了節省將 numpy 陣列複製到指定裝置,然後又將結果儲存到 numpy 陣列中所浪費的時間,Numba 提供了一些 函式 來宣告並將陣列送到指定裝置,如:numba.cuda.device_array
,numba.cuda。 device_array_like
,numba.cuda.to_device
等函式來節省不必要的複製到 cpu 的時間(除非必要)。
另一方面,裝置函式 只能從裝置內部(通過核函式或其他裝置函式)呼叫。 比較好的一點是,您可以從 裝置函式 中返回一個值。 因此,您可以用此函式的返回值來計算 核函式
或 裝置函式
裡的內容。
1 2 3 4 5 |
from numba import cuda @cuda.jit(device=True) def device_function(a, b): return a + b |
您還應該在這裡檢視 Numba 的 cuda 庫支援的功能。
Numba 在其 cuda 庫中也有自己的 原子操作,隨機數生成器,共享記憶體實現(以加快資料的訪問)等功能。
ctypes/cffi/cython 的互用性:
7. 擴充套件閱讀
- https://nbviewer.jupyter.org/github/ContinuumIO/gtc2017-numba/tree/master/
- https://devblogs.nvidia.com/seven-things-numba/
- https://devblogs.nvidia.com/numba-python-cuda-acceleration/
- https://jakevdp.github.io/blog/2015/02/24/optimizing-python-with-numpy-and-numba/
- https://www.youtube.com/watch?v=1AwG0T4gaO0
8. 參考
- http://numba.pydata.org/numba-doc/latest/user/index.html
- https://github.com/ContinuumIO/gtc2018-numba
- http://stephanhoyer.com/2015/04/09/numba-vs-cython-how-to-choose/
謝謝閱讀!
原文連結:https://towardsdatascience.com/speed-up-your-algorithms-part-2-numba-293e554c5cc1