如果你在使用Python進(jìn)行高性能計(jì)算，Numba提供的加速效果可以比肩原生的C/C++程序，只需要在函數(shù)上添加一行@jit的裝飾。它支持CPU和GPU，是數(shù)據(jù)科學(xué)家必不可少的編程利器

之前的文章計(jì)算機(jī)基礎(chǔ)系列：源代碼如何被計(jì)算機(jī)執(zhí)行已經(jīng)提到計(jì)算機(jī)只能執(zhí)行二進(jìn)制的機(jī)器碼，C、C++等編譯型語言依靠編譯器將源代碼轉(zhuǎn)化為可執(zhí)行文件后才能運(yùn)行，Python、Java等解釋型語言使用解釋器將源代碼翻譯后在虛擬機(jī)上執(zhí)行。對(duì)于Python，由于解釋器的存在，其執(zhí)行效率比C語言慢幾倍甚至幾十倍。

以C語言為基準(zhǔn)，不同編程語言性能測(cè)試比較

上圖比較了當(dāng)前流行的各大編程語言在幾個(gè)不同任務(wù)上的計(jì)算速度。C語言經(jīng)過幾十年的發(fā)展，優(yōu)化已經(jīng)達(dá)到了極致。以C語言為基準(zhǔn)，大多數(shù)解釋語言，如Python、R會(huì)慢十倍甚至一百倍。Julia這個(gè)解釋語言是個(gè)“奇葩”，因?yàn)樗捎昧薐IT編譯技術(shù)。

解決Python執(zhí)行效率低的問題，一種解決辦法是使用C/C++語言重寫Python函數(shù)，但是這要求程序員對(duì)C/C++語言熟悉，且調(diào)試速度慢，不適合絕大多數(shù)Python程序員。另外一種非常方便快捷的解決辦法就是使用Just-In-Time（JIT）技術(shù)，本文將解釋JIT技術(shù)的原理，并提供幾個(gè)案例，讓你十分鐘內(nèi)學(xué)會(huì)JIT技術(shù)。

Python解釋器工作原理

Python是一門解釋語言，Python為我們提供了基于硬件和操作系統(tǒng)的一個(gè)虛擬機(jī)，并使用解釋器將源代碼轉(zhuǎn)化為虛擬機(jī)可執(zhí)行的字節(jié)碼。字節(jié)碼在虛擬機(jī)上執(zhí)行，得到結(jié)果。

Python解釋器工作原理

我們使用python example.py來執(zhí)行一份源代碼時(shí)，Python解釋器會(huì)在后臺(tái)啟動(dòng)一個(gè)字節(jié)碼編譯器（Bytecode Compiler），將源代碼轉(zhuǎn)換為字節(jié)碼。字節(jié)碼是一種只能運(yùn)行在虛擬機(jī)上的文件，Python的字節(jié)碼默認(rèn)后綴為.pyc，Python生成.pyc后一般放在內(nèi)存中繼續(xù)使用，并不是每次都將.pyc文件保存到磁盤上。有時(shí)候我們會(huì)看到自己Python代碼文件夾里有很多.pyc文件與.py文件同名，但也有很多時(shí)候看不到.pyc文件。pyc字節(jié)碼通過Python虛擬機(jī)與硬件交互。虛擬機(jī)的出現(xiàn)導(dǎo)致程序和硬件之間增加了中間層，運(yùn)行效率大打折扣。相信使用過虛擬機(jī)軟件的朋友深有體會(huì)，在原生的系統(tǒng)上安裝一個(gè)虛擬機(jī)軟件，在虛擬機(jī)上再運(yùn)行一個(gè)其他系統(tǒng)，經(jīng)常感覺速度下降，體驗(yàn)變差，這與Python虛擬機(jī)導(dǎo)致程序運(yùn)行慢是一個(gè)原理。

Just-In-Time（JIT）技術(shù)為解釋語言提供了一種優(yōu)化，它能克服上述效率問題，極大提升代碼執(zhí)行速度，同時(shí)保留Python語言的易用性。使用JIT技術(shù)時(shí)，JIT編譯器將Python源代碼編譯成機(jī)器直接可以執(zhí)行的機(jī)器語言，并可以直接在CPU等硬件上運(yùn)行。這樣就跳過了原來的虛擬機(jī)，執(zhí)行速度幾乎與用C語言編程速度并無二致。

十分鐘上手Numba

Numba是一個(gè)針對(duì)Python的開源JIT編譯器，由Anaconda公司主導(dǎo)開發(fā)，可以對(duì)原生代碼進(jìn)行CPU和GPU加速。

使用conda安裝Numba：

$ conda install numba

或者使用pip安裝：

$ pip install numba

使用時(shí)，只需要在原來的函數(shù)上添加一行'注釋'：

from numba import jit
import numpy as np
SIZE = 2000
x = np.random.random((SIZE, SIZE))
'''
給定n*n矩陣，對(duì)矩陣每個(gè)元素計(jì)算tanh值，然后求和。
因?yàn)橐h(huán)矩陣中的每個(gè)元素，計(jì)算復(fù)雜度為 n*n。
'''
@jit
def jit_tan_sum(a): # 函數(shù)在被調(diào)用時(shí)編譯成機(jī)器語言
tan_sum = 0
for i in range(SIZE): # Numba 支持循環(huán)
for j in range(SIZE):
tan_sum += np.tanh(a[i, j]) # Numba 支持絕大多數(shù)NumPy函數(shù)
return tan_sum
print(jit_tan_sum(x))

我們只需要在原來的代碼上添加一行@jit，即可將一個(gè)函數(shù)編譯成機(jī)器碼，其他地方都不需要更改。@符號(hào)裝飾了原來的代碼，所以稱類似寫法為裝飾器。

在我的Core i5處理器上，添加@jit裝飾器后，上面的代碼執(zhí)行速度提升了23倍！而且隨著數(shù)據(jù)和計(jì)算量的增大，numba的性能提升可能會(huì)更大！很多朋友的代碼可能需要執(zhí)行十幾個(gè)小時(shí)甚至一天，進(jìn)行加速，完全有可能把一天的計(jì)算量縮短到幾個(gè)小時(shí)！

Numba的使用場(chǎng)景

Numba簡(jiǎn)單到只需要在函數(shù)上加一個(gè)裝飾就能加速程序，但也有缺點(diǎn)。目前Numba只支持了Python原生函數(shù)和部分NumPy函數(shù)，其他一些場(chǎng)景可能不適用。比如Numba官方給出這樣的例子：

from numba import jit
import pandas as pd
x = {'a': [1, 2, 3], 'b': [20, 30, 40]}
@jit
def use_pandas(a): # Function will not benefit from Numba jit
df = pd.DataFrame.from_dict(a) # Numba doesn't know about pd.DataFrame
df += 1 # Numba doesn't understand what this is
return df.cov() # or this!
print(use_pandas(x))

pandas是更高層次的封裝，Numba其實(shí)不能理解它里面做了什么，所以無法對(duì)其加速。一些大家經(jīng)常用的機(jī)器學(xué)習(xí)框架，如scikit-learn，tensorflow，pytorch等，已經(jīng)做了大量的優(yōu)化，不適合再使用Numba做加速。

此外，Numba不支持：

• try…except 異常處理
• with 語句
• yield from

注 Numba當(dāng)前支持的功能： http://numba./numba-doc/latest/reference/pysupported.html

那如何決定是否使用Numba呢？

Numba的@jit裝飾器就像自動(dòng)駕駛，用戶不需要關(guān)注到底是如何優(yōu)化的，Numba去嘗試進(jìn)行優(yōu)化，如果發(fā)現(xiàn)不支持，那么Numba會(huì)繼續(xù)用Python原來的方法去執(zhí)行該函數(shù)，即圖 Python解釋器工作原理中左側(cè)部分。這種模式被稱為object模式。前文提到的pandas的例子，Numba發(fā)現(xiàn)無法理解里面的內(nèi)容，于是自動(dòng)進(jìn)入了object模式。object模式還是和原生的Python一樣慢，還有可能比原來更慢。

Numba真正牛逼之處在于其nopython模式。將裝飾器改為@jit(nopython=True)或者@njit，Numba會(huì)假設(shè)你已經(jīng)對(duì)所加速的函數(shù)非常了解，強(qiáng)制使用加速的方式，不會(huì)進(jìn)入object模式，如編譯不成功，則直接拋出異常。nopython的名字會(huì)有點(diǎn)歧義，我們可以理解為不使用很慢的Python，強(qiáng)制進(jìn)入圖 Python解釋器工作原理中右側(cè)部分。

實(shí)踐上，一般推薦將代碼中計(jì)算密集的部分作為單獨(dú)的函數(shù)提出來，并使用nopython方式優(yōu)化，這樣可以保證我們能使用到Numba的加速功能。其余部分還是使用Python原生代碼，在計(jì)算加速的前提下，避免過長(zhǎng)的編譯時(shí)間。（有關(guān)編譯時(shí)間的問題下節(jié)將會(huì)介紹。）Numba可以與NumPy緊密結(jié)合，兩者一起，常常能夠得到近乎C語言的速度。盡管Numba不能直接優(yōu)化pandas，但是我們可以將pandas中處理數(shù)據(jù)的for循環(huán)作為單獨(dú)的函數(shù)提出來，再使用Numba加速。

編譯開銷

編譯源代碼需要一定的時(shí)間。C/C++等編譯型語言要提前把整個(gè)程序先編譯好，再執(zhí)行可執(zhí)行文件。Numba庫提供的是一種懶編譯（Lazy Compilation）技術(shù)，即在運(yùn)行過程中第一次發(fā)現(xiàn)代碼中有@jit，才將該代碼塊編譯。用到的時(shí)候才編譯，看起來比較懶，所以叫懶編譯。使用Numba時(shí)，總時(shí)間 = 編譯時(shí)間 + 運(yùn)行時(shí)間。相比所能節(jié)省的計(jì)算時(shí)間，編譯的時(shí)間開銷很小，所以物有所值。對(duì)于一個(gè)需要多次調(diào)用的Numba函數(shù)，只需要編譯一次，后面再調(diào)用時(shí)就不需要編譯了。

from numba import jit
import numpy as np
import time
SIZE = 2000
x = np.random.random((SIZE, SIZE))
'''
給定n*n矩陣，對(duì)矩陣每個(gè)元素計(jì)算tanh值，然后求和。
因?yàn)橐h(huán)矩陣中的每個(gè)元素，計(jì)算復(fù)雜度為 n*n。
'''
@jit
def jit_tan_sum(a): # 函數(shù)在被調(diào)用時(shí)編譯成機(jī)器語言
tan_sum = 0
for i in range(SIZE): # Numba 支持循環(huán)
for j in range(SIZE):
tan_sum += np.tanh(a[i, j]) # Numba 支持絕大多數(shù)NumPy函數(shù)
return tan_sum
# 總時(shí)間 = 編譯時(shí)間 + 運(yùn)行時(shí)間
start = time.time()
jit_tan_sum(x)
end = time.time()
print('Elapsed (with compilation) = %s' % (end - start))
# Numba將加速的代碼緩存下來
# 總時(shí)間 = 運(yùn)行時(shí)間
start = time.time()
jit_tan_sum(x)
end = time.time()
print('Elapsed (after compilation) = %s' % (end - start))

代碼中兩次調(diào)用Numba優(yōu)化函數(shù)，第一次執(zhí)行時(shí)需要編譯，第二次使用緩存的代碼，運(yùn)行時(shí)間將大大縮短：

Elapsed (with compilation) = 0.49199914932250977
Elapsed (after compilation) = 0.0364077091217041

原生Python速度慢的另一個(gè)重要原因是變量類型不確定。聲明一個(gè)變量的語法很簡(jiǎn)單，如a = 1，但沒有指定a到底是一個(gè)整數(shù)和一個(gè)浮點(diǎn)小數(shù)。Python解釋器要進(jìn)行大量的類型推斷，會(huì)非常耗時(shí)。同樣，引入Numba后，Numba也要推斷輸入輸出的類型，才能轉(zhuǎn)化為機(jī)器碼。針對(duì)這個(gè)問題，Numba給出了名為Eager Compilation的優(yōu)化方式。

from numba import jit, int32
@jit('int32(int32, int32)', nopython=True)
def f2(x, y):
# A somewhat trivial example
return x + y

@jit(int32(int32, int32))告知Numba你的函數(shù)在使用什么樣的輸入和輸出，括號(hào)內(nèi)是輸入，括號(hào)左側(cè)是輸出。這樣不會(huì)加快執(zhí)行速度，但是會(huì)加快編譯速度，可以更快將函數(shù)編譯到機(jī)器碼上。

Numba到底有多快

網(wǎng)上有很多對(duì)Numba進(jìn)行性能評(píng)測(cè)的文章，在一些計(jì)算任務(wù)上，Numba結(jié)合NumPy，可得到接近C語言的速度。

Numba性能測(cè)試

Numba的更多功能

除了上面介紹的加速功能，Numba還有很多其他功能。@vectorize裝飾器可以將一個(gè)函數(shù)向量化，變成類似NumPy函數(shù)一樣，直接處理矩陣和張量。R語言用戶可能非常喜歡這個(gè)功能。

Numba還可以使用GPU進(jìn)行加速，目前支持英偉達(dá)的CUDA和AMD的ROC。GPU的工作原理和編程方法與CPU略有不同，本專欄會(huì)在后續(xù)文章中介紹GPU編程。

Numba原理

Numba編譯過程

Numba使用了LLVM和NVVM技術(shù)，這個(gè)技術(shù)可以將Python、Julia這樣的解釋語言直接翻譯成CPU或GPU可執(zhí)行的機(jī)器碼。

小結(jié)

無論你是在做金融量化分析，還是計(jì)算機(jī)視覺，如果你在使用Python進(jìn)行高性能計(jì)算，處理矩陣和張量，或包含其他計(jì)算密集型運(yùn)算，Numba提供的加速效果可以比肩原生的C/C++程序，只需要在函數(shù)上添加一行@jit的裝飾。它支持CPU和GPU，是數(shù)據(jù)科學(xué)家必不可少的編程利器。