Python Numba CPU下加速

Python代碼加速

主要考慮代碼優(yōu)化加速，而非代碼邏輯優(yōu)化。

Python代碼直接運(yùn)行GPU是不行的，需要一定的改變，Numba是一個(gè)接口，不過本文主要針對CPU下的Python代碼加速。

Python解釋器工作原理

Python文件執(zhí)行過程

.py文件通過解釋器轉(zhuǎn)化為虛擬機(jī)可以執(zhí)行的字節(jié)碼(.pyc)；

字節(jié)碼在虛擬機(jī)上執(zhí)行，得到結(jié)果；

字節(jié)碼是一種只能運(yùn)行在虛擬機(jī)上的文件，默認(rèn)后綴.pyc，Python生成.pyc之后一般放在內(nèi)存中繼續(xù)使用，并不是每次都將.pyc文件保存到磁盤上。

虛擬機(jī)是基于硬件和操作系統(tǒng)的。

.pyc字節(jié)碼通過Python虛擬機(jī)與硬件交互，就是說虛擬機(jī)的存在導(dǎo)致程序與硬件之間增加了中間層。效率自然被拖后。

JIT（Just-In-Time）

JIT技術(shù)中，JIT編譯器將Python源代碼.py直接編譯成機(jī)器可以執(zhí)行的機(jī)器語言（機(jī)器碼），就可以直接在CPU等硬件上運(yùn)行。

這樣，JIT就跳過了原來的虛擬機(jī)，執(zhí)行速度幾乎與用C語言編程速度無差別。

Numba庫

Numba是Anaconda公司開發(fā)的針對Python的開源JIT編譯器，用于提供Python版CPU和GPU編程，速度比原生Python快數(shù)十倍。一、安裝Numba

pip install numba

# 或者

conda install numba

二、使用方法：裝飾器

from numba import jit

import numpy as np

SIZE = 2000

x = np.random.random((SIZE, SIZE))

"""

給定n*n矩陣，對矩陣每個(gè)元素計(jì)算tanh值，然后求和。

因?yàn)橐h(huán)矩陣中的每個(gè)元素，計(jì)算復(fù)雜度為 n*n。

"""

@jit # numba的使用方法

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

如代碼所示，只是在函數(shù)上加一個(gè)裝飾器@jit，就可以將運(yùn)行速度提升20多倍。

@jit裝飾器的本質(zhì)，是將函數(shù)編譯為機(jī)器碼，省掉虛擬機(jī)環(huán)節(jié)，提升速度。

三、Numba的使用場景總結(jié)

numba目前只支持Python原生函數(shù)和部分Numpy函數(shù)，其他場景下無效。

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，而Pandas并不是原生代碼，而是更高層次的封裝，Numba不能理解pandas內(nèi)部在做什么，所以無法對其加速。

而一些常用的機(jī)器學(xué)習(xí)框架，比如scikit-learn, tensorflow, pyrorch等，已經(jīng)做了大量的優(yōu)化，不適合再使用Numba做加速。

可以簡單總結(jié)為，Numba不支持：

pandas

scikit-learn, tensorflow, pyrorch

try…except 異常處理

with 語句

yield from

四、Numba具體使用過程

Numba有兩種模式：

@jit：object模式：上圖左側(cè)

Numba的@jit裝飾器會嘗試優(yōu)化代碼，如果發(fā)現(xiàn)不支持（比如pandas等），那么Numba會繼續(xù)使用Python原來的方法去執(zhí)行該函數(shù)。

@jit(nopython=True)或者@njit：nopython模式：上圖右側(cè)

強(qiáng)制加速，不會進(jìn)入上圖左側(cè)流程，只進(jìn)行右側(cè)流程，如果編譯不成功，就拋出異常。

實(shí)際使用中，一把推薦將代碼中計(jì)算密集的部分作為單獨(dú)的函數(shù)提出來，并使用nopython方法優(yōu)化，這樣可以保證能用到Numba加速；其余部分還是使用Python原生代碼。

五、編譯開銷

編譯源代碼需要一定的時(shí)間：

C/C++等編譯型語言是提前把整個(gè)程序先編譯好，再執(zhí)行可執(zhí)行文件；

Numba庫是懶編譯（Lazy Compilation）技術(shù)；

其中，懶編譯技術(shù)（Lazy Compilation）即

在運(yùn)行過程中第一次發(fā)現(xiàn)源代碼中有@jit，才將該代碼塊編譯；

同一個(gè)Numba函數(shù)多次調(diào)用，只需要編譯一次；

總 時(shí) 間 = 編 譯 時(shí) 間 + 運(yùn) 行 時(shí) 間 總時(shí)間=編譯時(shí)間+運(yùn)行時(shí)間

總時(shí)間=編譯時(shí)間+運(yùn)行時(shí)間

from numba import jit

import numpy as np

import time

SIZE = 2000

x = np.random.random((SIZE, SIZE))

"""

給定n*n矩陣，對矩陣每個(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í)需要編譯，第二次就直接使用緩存的已經(jīng)編譯好的代碼，運(yùn)行時(shí)間大大縮短。

六、確定輸入輸出類型Eager Compilation：節(jié)省編譯速度

原生Python速度慢的另一個(gè)因素是變量類型不確定，Python解釋器需要進(jìn)行大量的類型推斷。

Numba也要推斷輸入輸出的類型：

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ù)在使用什么樣的輸入和輸出：

括號內(nèi)是輸入；

括號左側(cè)是輸出；

這樣不會加快執(zhí)行速度，但是會加快編譯速度，可以更快將函數(shù)編譯到機(jī)器碼上。

Numba原理

Numba使用了LLVM和NVVM技術(shù)，此技術(shù)將Python等解釋型語言直接翻譯成CPU、GPU可執(zhí)行的機(jī)器碼。

Python 虛擬化

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。

版權(quán)聲明：本文內(nèi)容由網(wǎng)絡(luò)用戶投稿，版權(quán)歸原作者所有，本站不擁有其著作權(quán)，亦不承擔(dān)相應(yīng)法律責(zé)任。如果您發(fā)現(xiàn)本站中有涉嫌抄襲或描述失實(shí)的內(nèi)容，請聯(lián)系我們jiasou666@gmail.com 處理，核實(shí)后本網(wǎng)站將在24小時(shí)內(nèi)刪除侵權(quán)內(nèi)容。