Python高性能計算庫——Numba

最近我在觀看一些SciPy2017會議的視頻，偶然發現關於Numba的來歷--講述了那些C++的高手們因為對Gil Forsyth和Lorena Barba失去信心而編寫的一個庫。雖然本人覺得這個做法有些不妥，但我真的很喜歡他們所分享的知識。因為我發現自己正在受益於這個庫，並且從Python代碼中獲得了令人難以置信的表現，所以我覺得應該要寫一些關於Numba庫的介紹性文章，也可能會在將來添加一系列小的更多類似教程的文章。

1.那麼到底什麼是Numba？

Numba是一個庫，可以在運行時將Python代碼編譯為本地機器指令，而不會強制大幅度的改變普通的Python代碼（稍後再做說明）。翻譯/魔術是使用LLVM編譯器完成的，該編譯器是相當活躍的開源社區開發的。

Numba最初是由Continuum Analytics內部開發，此公司也開發了著名的Anaconda，但現在它是開源的。核心應用領域是math-heavy（密集數學？重型數學？）和array-oriented（面向數組）的功能，它們在本地Python中相當緩慢。想像一下，在Python中編寫一個模塊，必須一個元素接著一個元素的循環遍歷一個非常大的數組來執行一些計算，而不能使用向量操作來重寫。這是很不好的主意，是吧？所以「通常」這類庫函數是用C / C ++或Fortran編寫的，編譯後，在Python中作為外部庫使用。Numba這類函數也可以寫在普通的Python模塊中，而且運行速度的差別正在逐漸縮小。

2.怎麼才能get到Numba呢？

安裝Numba的推薦方法是使用conda包管理

conda install numba

你也可以用pip來安裝Numba，但是最新版本的發布才一天之久。但是，只要你能夠使用conda，我會推薦使用它，因為它能夠為你安裝例如CUDA工具包，也許你想讓你的Python代碼GPU就緒（當然，這也是有可能的！）。

3.如何使用Numba呢？

使用它的要求不多。基本上，你寫一個自己的「普通」的Python函數，然後給函數定義添加一個裝飾（如果你不是很熟悉裝飾器，讀一下關於this或that）。你可以使用不同類型的裝飾器，但@jit可能是剛開始的選擇之一。其他裝飾器可用於例如創建numpy通用功能@vectorize或編寫將在CUDA GPU上執行的代碼@cuda。我不會在這篇文章中介紹這些裝飾。現在，讓我們來看看基本的步驟。他們提供的代碼示例是2d數組的求和函數，以下是代碼：

fromnumbaimportjit

fromnumpyimportarange

# jit decorator tells Numba to compile this function.# The argument types will be inferred by Numba when function is called.

@jit

defsum2d(arr):

M,N=arr.shape

result=0.0

foriinrange(M):

forjinrange(N):

result+=arr[i,j]

returnresult

a=arange(9).reshape(3,3)

print(sum2d(a))

正如你所看到的，Numba裝飾器被添加到函數定義中，並且voilá這個函數將運行得很快。但是，這裡帶來了很有趣的注意事項：你只能使用Numpy和標準庫里的函數來加快Numba速度，甚至不需要開了他們所有的特性。他們有一個相當好的文檔（參考資料），列出了所有支持的內容。見here是所支持Python的功能和here 是所支持的Numpy功能。現在支持的功能可能還不太多，但我想告訴你，這就夠了！請記住，Numba不是要加快你的資料庫查詢或如何強化圖像處理功能。他們的目標是加快面向數組的計算，我們可以使用它們庫中提供的函數來解決。

4.示例和速度比較

熟練的Python用戶永遠不會使用上述代碼實現sum功能，而是調用numpy.sum。相反，我將向你介紹另外一個例子，為了更好地理解這個例子，也許剛開始是一個小的背景故事（如果你對這個例子的背景不感興趣，你可以直接跳過然後直接去看代碼）。

從我所學習的知識來看，我會認為自己是一個水文學家，我做的很多的一件事是模擬降雨徑流過程。簡單點來說：通過時間序列數據，例如雨量和空氣溫度，然後嘗試創建模型來判斷一條河流的水流量有多少。這在外行看來是非常複雜。但，對於我們來說，很簡單。我們通常使用的模塊迭代輸入數組，並且對於每個時間步長，我們會更新一些模塊內部的狀態（例如，模擬土壤水分，積雪或攔截水中的樹木）。在每個時間段結束時，計算水流量，這不僅取決於在同一時間步長下的雨，而且也取決於在內部模型狀態（或儲存）。在這種情況下，我們就需要考慮以前時間步長的狀態和輸出。那麼你可能會看到這個問題：我們必須一段時間接一段時間的計算整個流程，而對於解決這種問題Python本來就是很慢的！這就是為什麼大多數模塊都是在Fortran或C/C ++中實現的。如前所述：Python在對於這種面向數組的計算來說是慢的。但是Numba允許我們在Python中做同樣的事情，而且沒有太多的性能損失。我認為至少對於模型的理解和發展，這可能會很方便。（所以我最近創建了一個名為「RRMPG 」的項目——降雨徑流建模遊樂場）。

Okay，現在我們來看看我們get到了什麼。我們將使用最簡單的模塊之一，由MB Fiering在1967年出於教育目的開發的ABC模型，並將Python代碼的速度與Numba優化後Python代碼和Fortran實現進行比較。請注意這個模型不是我們在現實中使用的（正如名稱所示），但是我認為這可能是一個不錯的想法來舉例。

A、B、C模塊是一個三個參數模塊（a，b，c，習慣性命名），它只接收下雨量為輸入，只有一個存儲。土壤水分蒸發蒸騰損失總量（參數b），另一部分通過土壤滲透到地下水儲存（參數a），最後一個參數c代表地下水總量，離開地下變成河流。Python中的代碼，使用Numpy數組可能會像如下所示：

importnumpyasnp

defabc_model_py(a,b,c,rain):

# initialize array for the stream discharge of each time step

outflow=np.zeros((rain.size),dtype=np.float64)

# placeholder, in which we save the storage content of the previous and

# current timestep

state_in=0

state_out=0

foriinrange(rain.size):

# Update the storage

state_out=(1-c)*state_in+a*rain[i]

# Calculate the stream discharge

outflow[i]=(1-a-b)*rain[i]+c*state_out

state_in=state_out

returnoutflow

接下來我們使用Numba來實現相同的功能。

@jit

defabc_model_numba(a,b,c,rain):

outflow=np.zeros((rain.size),dtype=np.float64)

state_in=0

state_out=0

foriinrange(rain.size):

state_out=(1-c)*state_in+a*rain[i]

outflow[i]=(1-a-b)*rain[i]+c*state_out

state_in=state_out

returnoutflow

我用隨機數字作為輸入來運行這些模塊，這只是為了比較計算時間，而且也比較了針對fortran實現的時間（詳見here）。我們來看看數字：

py_time =%timeit -r 5-n 10-o abc_model_py(0.2, 0.6, 0.1, rain)

>>6.75 s ±11.6 ms per loop (mean ± std. dev. of 5 runs, 10 loops each)

# Measure the execution time of the Numba implementation

numba_time =%timeit -r 5-n 10-o abc_model_numba(0.2, 0.6, 0.1, rain)

>>30.6 ms ±498μs per loop (mean ± std. dev. of 5 runs, 10 loops each)

# Measure the execution time of the Fortran implementation

fortran_time =%timeit -r 5-n 10-o abc_model_fortran(0.2, 0.6, 0.1, rain)

>>31.9 ms ±757μs per loop (mean ± std. dev. of 5 runs, 10 loops each)

# Compare the pure Python vs Numba optimized time

py_time.best / numba_time.best

>>222.1521754580626

# Compare the time of the fastes numba and fortran run

numba_time.best / fortran_time.best

>>0.9627960721576471

通過添加一個裝飾器，我們的計算速度比純Python代碼快222倍，甚至比Fortran也快很多。在計算能力決定未來的時代，Numba一定會被更多人接受。

以上就是我的介紹，希望有人現在有動力去看看Numba庫。我想在將來我會編寫一系列小的Numba文章/教程，並提供更多的技術信息，讓更多的人使用Numba 庫。而本文僅作為一個開始。

本文由北郵@愛可可-愛生活 老師推薦，阿里云云棲社區組織翻譯。

文章原標題《Introduction to the Numba library》

作者：Flaire

譯者：一隻高冷的貓，審校：袁虎。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！