技術背景
python作為一門編程語言，有非常大的生態優勢，但是其執行效率一直被人詬病。純粹的python代碼跑起來速度會非常的緩慢，因此很多對性能要求比較高的python庫，需要用C++或者Fortran來構造底層算法模塊，再用python進行上層封裝的方案。在前面寫過的這篇博客中，介紹了使用f2py將fortran代碼編譯成動態鏈接庫的方案，這可以認為是一種“事前編譯”的手段。但是本文將要介紹一種即時編譯（Just In Time，簡稱JIT）的手段，也就是在臨近執行函數前，才對其進行編譯。以下截圖來自於參考鏈接4，講述了關於常見的一些編譯場景的區別：
用numba.jit加速求平方和
numba中大部分加速的函數都是通過裝飾器（decorator）來實現的，關於python中decorator的使用方法和場景，在前面寫過的這篇博客中有比較詳細的介紹，讓我們直接使用numba的裝飾器來解決一些實際問題。這里的問題場景是，隨便給定一個數列，在不用求和公式的情況下對這個數列的所有元素求平方和，即：
我們已知類似於這種求和的形式，其實是有很大的優化空間的，相比於直接用一個for循環來求解的話。這里我們直接展示一下案例代碼：
# test_jit.py
from numba import jit
import time
import matplotlib.pyplot as plt
def adder(max): # 普通的循環求解
s = 0
for i in range(max):
s += i ** 2
return s
@jit(nopython=True)
def jit_adder(max): # 使用即時編譯求解
s = 0
for i in range(max):
s += i ** 2
return s
if __name__ == '__main__':
time_adder = []
time_jit_adder = []
x = list(range(1, 10000000, 500000))
for i in x:
time1 = time.time()
s = adder(i)
time2 = time.time()
s = jit_adder(i)
time3 = time.time()
time_adder.append(time2 - time1)
time_jit_adder.append(time3 - time2)
# 開始作圖
fig, ax1 = plt.subplots()
color = 'black'
ax1.set_xlabel('Numbers')
ax1.set_ylabel('Time (s)', color=color)
ax1.plot(x[1:], time_adder[1:], color=color, label='python')
ax1.tick_params(axis='y', labelcolor=color)
ax2 = ax1.twinx() # 第二個y-坐標軸
color = 'red'
ax2.set_ylabel('Time (s)', color=color)
ax2.plot(x[1:], time_jit_adder[1:], color=color, label='jit')
ax2.tick_params(axis='y', labelcolor=color)
plt.title('Running time difference via using jit')
fig.tight_layout()
plt.legend()
plt.savefig('jit.png')
運行該python文件，會在當前目錄下產生一個雙坐標軸的圖像：
在這個計算結果中，使用了即時編譯技術之后，求解的時間幾乎被壓縮到了微秒級別，而循環求和的方法卻已經達到了秒級，加速倍數在\(10^5\)級別。
用numba.jit加速求雙曲正切函數和
在上一個案例中，也許涉及到的計算過於的簡單，導致了加速倍數超出了想象的情況。因此這里我們只替換所求解的函數，看看加速的倍數是否會發生變化。這里我們采用了雙曲正切求和的函數：
通過math來實現這個函數的計算，用以替換上一章節中求平方值的方法：
# test_jit.py
from numba import jit
import time
import matplotlib.pyplot as plt
import math
def adder(max):
s = 0
for i in range(max):
s += math.tanh(i ** 2)
return s
@jit(nopython=True)
def jit_adder(max):
s = 0
for i in range(max):
s += math.tanh(i ** 2)
return s
if __name__ == '__main__':
time_adder = []
time_jit_adder = []
x = list(range(1, 10000000, 500000))
for i in x:
time1 = time.time()
s = adder(i)
time2 = time.time()
s = jit_adder(i)
time3 = time.time()
time_adder.append(time2 - time1)
time_jit_adder.append(time3 - time2)
fig, ax1 = plt.subplots()
color = 'black'
ax1.set_xlabel('Numbers')
ax1.set_ylabel('Time (s)', color=color)
ax1.plot(x[1:], time_adder[1:], color=color, label='python')
ax1.tick_params(axis='y', labelcolor=color)
ax2 = ax1.twinx()
color = 'red'
ax2.set_ylabel('Time (s)', color=color)
ax2.plot(x[1:], time_jit_adder[1:], color=color, label='jit')
ax2.tick_params(axis='y', labelcolor=color)
plt.title('Running time difference via using jit')
fig.tight_layout()
plt.legend()
plt.savefig('jit.png')
最終得到的時間對比圖結果如下所示：
需要提醒的是，黑色的曲線所對應的坐標軸是左邊黑色標識的坐標軸，而紅色的曲線所對應的坐標軸是右邊紅色標識的坐標軸。因此，這個圖給我們的提示信息是，使用即時編譯技術之后，加速的倍率大約為\(10^2\)。這個加速倍率相對來說更加可以接受，因為C++等語言比python直接計算的速度在特定場景下大概就是要快上幾百倍。
用numba.vectorize執行向量化計算
關於向量化計算的原理和方法，在這篇文章中有比較好的描述，這里放上部分截圖說明：
總結為，向量化計算的方法本質上也是一種並行化計算的方法，並行化技術的可行性是來源於SIMD技術，在指令集的層面對數據進行並行化的處理。在numpy的庫中是自帶支持SIMD的向量化計算的，因此速度非常的高，比如numpy.dot函數就是通過向量化計算來實現的。但是numpy能夠執行的任務僅僅局限在numpy自身所支持的有限的函數上，因此如果是需要一個不同的函數，那么就需要用到numba的向量化計算模塊了。
# test_vectorize.py
from numba import vectorize
import numpy as np
import time
import matplotlib.pyplot as plt
def ddot(max):
s = 0
np.random.seed(1)
a1 = np.random.randn(max)
np.random.seed(2)
a2 = np.random.randn(max)
for i in range(max):
s += a1[i] * a2[i]
return s
@vectorize
def jit_ddot(max):
s = 0
np.random.seed(1)
a1 = np.random.randn(max)
np.random.seed(2)
a2 = np.random.randn(max)
for i in range(max):
s += a1[i] * a2[i]
return s
def numpy_ddot(max):
np.random.seed(1)
a1 = np.random.randn(max)
np.random.seed(2)
a2 = np.random.randn(max)
return np.dot(a1, a2)
if __name__ == '__main__':
time_ddot = []
time_jit_ddot = []
time_numpy_ddot = []
x = list(range(1, 1000000, 50000))
for i in x:
time1 = time.time()
s = ddot(i)
time2 = time.time()
s = jit_ddot(i)
time3 = time.time()
s = numpy_ddot(i)
time4 = time.time()
time_ddot.append(time2 - time1)
time_jit_ddot.append(time3 - time2)
time_numpy_ddot.append(time4 - time3)
fig, ax1 = plt.subplots()
color = 'black'
ax1.set_xlabel('Numbers')
ax1.set_ylabel('Time (s)', color=color)
ax1.plot(x[1:], time_ddot[1:], color=color, label='python')
ax1.tick_params(axis='y', labelcolor=color)
ax2 = ax1.twinx()
color = 'red'
ax2.set_ylabel('Time (s)', color=color)
ax2.plot(x[1:], time_jit_ddot[1:], color=color, label='jit')
ax2.plot(x[1:], time_numpy_ddot[1:], 's', color=color, label='numpy')
ax2.tick_params(axis='y', labelcolor=color)
plt.title('Running time difference via using jit')
fig.tight_layout()
plt.legend()
plt.savefig('jit.png')
運行結果如下：
可以看到雖然相比與numpy的同樣的向量化計算方法，numba速度略慢一些，但是都比純粹的python代碼性能要高兩個量級。這里也給我們一個啟發，如果追求極致的性能，最好是盡可能的使用numpy中已有的函數。當然，在一些數學函數的計算上，numpy的速度比math還是要慢上一些的，這里我們就不展開介紹了。
總結概要
本文介紹了numba的兩個裝飾器的原理與測試案例，以及python中兩坐標軸繪圖的案例。其中基於即時編譯技術jit的裝飾器，能夠對代碼中的for循環產生較大的編譯優化，可以配合並行技術使用。而基於SIMD的向量化計算技術，也能夠在向量的計算中，如向量間的乘加運算等場景中，實現巨大的加速效果。這都是非常底層的優化技術，但是要分場景使用，numba這個強力的工具並不能保證在所有的計算場景下都能夠產生如此的加速效果。
版權聲明
本文首發鏈接為：https://www.cnblogs.com/dechinphy/p/numba.html
作者ID：DechinPhy
更多原著文章請參考：https://www.cnblogs.com/dechinphy/