我們都知道 Python 比較慢，但很多時候我們都不知道為什麼。雖然我用 Python 也有那麼兩年左右了，但也只能模模糊糊地感受到這麼兩點：
* Python 太動態了
* 如果能事先編譯一下 Python，讓它靜態一點，速度應該就會上來
於是我們就有了 cython。然而 cython 畢竟不是原生的 Python 代碼，使用起來還是有諸多不便的。為此，numba 就成了一個功能強大又容易上手的替代選擇。下面我們就先來看一下它的基本用法，在最後我們則會用卷積神經網絡（CNN）的卷積和池化操作來直觀感受一下其威力
（注意：以下代碼的運行環境均為 Jupyter Notebook、Python3.6.1）
使用 jit 加速 Python 低效的 for 語句jit 的全稱是 Just-in-time，在 numba 裡面則特指 Just-in-time compilation（即時編譯）。它背後的原理我們就不細說了，總之我們看到「編譯」兩個字大概就能感受到它是幹什麼的對吧（餵
那麼來看一個簡單的慄子——給數組中的每個數加上一個常數 c：
import numba as nb
import numpy as np
# 普通的 for
def add1(x, c):
rs = [0.] * len(x)
for i, xx in enumerate(x):
rs[i] = xx + c
return rs
# list comprehension
def add2(x, c):
return [xx + c for xx in x]
# 使用 jit 加速後的 for
@nb.jit(nopython=True)
def add_with_jit(x, c):
rs = [0.] * len(x)
for i, xx in enumerate(x):
rs[i] = xx + c
return rs
y = np.random.random(10**5).astype(np.float32)
x = y.tolist()
assert np.allclose(add1(x, 1), add2(x, 1), add_with_jit(x, 1))
%timeit add1(x, 1)
%timeit add2(x, 1)
%timeit add_with_jit(x, 1)
print(np.allclose(wrong_add(x, 1), 1))
以下是程序運行結果：
9.92 ms ± 188 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
5.77 ms ± 347 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.48 ms ± 171 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
需要注意的是：
* numba不支持 list comprehension，詳情可參見這裡
* jit能夠加速的不限於for，但一般而言加速for會比較常見、效果也比較顯著。我在我實現的numpy版本的卷積神經網絡（CNN）中用了jit後、可以把代碼加速 20 倍左右。具體代碼可以參見這裡，不過如果不想看原始碼的話，可以參見CNN.ipynb，我在其中做了一些相應的、比較簡單的實驗
* jit會在某種程度上「預編譯」你的代碼，這意味著它會在某種程度上固定住各個變量的數據類型；所以在jit下定義數組時，如果想要使用的是float數組的話，就不能用[0] * len(x)定義、而應該像上面那樣在0後面加一個小數點：[0.] * len(x)
雖然jit確實能讓我們代碼加速不少，但比之numpy的Ufunc還是要差很多：
assert np.allclose(y + 1, add_with_jit(x, 1))
%timeit add_with_jit(x, 1)
%timeit y + 1
結果將會是：
3.76 ms ± 292 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
19.8 µs ± 426 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
可以看到幾乎有 200 倍的差距，這當然是無法忍受的。為此，我們可以用vectorize來定義出類似於Ufunc的函數：
@nb.vectorize(nopython=True)
def add_with_vec(yy, c):
return yy + c
assert np.allclose(y + 1, add_with_vec(y, 1), add_with_vec(y, 1.))
%timeit add_with_vec(y, 1)
%timeit add_with_vec(y, 1.)
%timeit y + 1
%timeit y + 1.
上述程序的運行結果將會是：
72.5 µs ± 3.46 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
64.2 µs ± 1.9 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
24.6 µs ± 1.81 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
25.3 µs ± 1.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
雖然還是慢了 2 倍左右，但已經好很多了
然後有幾點需要注意的地方：
* vectorize下的函數所接受的參數都是一個個的數而非整個數組。所以上述add_with_vec的參數yy其實是輸入數組y中的元素，而不是y本身。更詳細的說明可以參見官方文檔）
* 可以看到當常數 c 是整數和是浮點數時、速度是不同的。個人猜測這是因為若常數 c 為整數，那麼實際運算時需要將它轉化為浮點數，從而導致速度變慢
* 上述代碼中我們沒有顯式地定義函數的參數類型和返回類型，但我們可以預先定義。比如說，如果我確定常數 c 就是整數的話，我就可以這樣寫：
@nb.vectorize("float32(float32, int32)", nopython=True)
def add_with_vec(yy, c):
return yy + c
而如果我確定常數 c 就是浮點數的話，我就可以這樣寫：
@nb.vectorize("float32(float32, float32)", nopython=True)
def add_with_vec(yy, c):
return yy + c
而如果我確定常數 c 不是整數就是浮點數的話（這個人好囉嗦！），我就可以這樣寫：
@nb.vectorize([
"float32(float32, int32)",
"float32(float32, float32)"
], nopython=True)
def add_with_vec(yy, c):
return yy + c
注意，float32 和 float64、int32 和 int64 是不同的，需要小心
此外，vectorize最炫酷的地方在於，它可以「並行」：
@nb.vectorize("float32(float32, float32)", target="parallel", nopython=True)
def add_with_vec(y, c):
return y + c
assert np.allclose(y+1, add_with_vec(y,1.))
%timeit add_with_vec(y, 1.)
%timeit y + 1
雖說在普通的 Python3.6.1 下、運行結果將如下：
73.5 µs ± 4.22 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
21.2 µs ± 734 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
似乎還變慢了；不過如果使用 Intel Distribution for Python 的話，會發現parallel版本甚至會比numpy原生的版本要稍快一些
那麼是否有用parallel總會更好的慄子呢？當然是有的：
# 將數組所有元素限制在某個區間[a, b]內
# 小於 a 則置為 a，大於 b 則置為 b
# 經典應用：ReLU
@nb.vectorize("float32(float32, float32, float32)", target="parallel", nopython=True)
def clip_with_parallel(y, a, b):
if y < a:
return a
if y > b:
return b
return y
@nb.vectorize("float32(float32, float32, float32)", nopython=True)
def clip(y, a, b):
if y < a:
return a
if y > b:
return b
return y
assert np.allclose(np.clip(y, 0.1, 0.9), clip(y, 0.1, 0.9), clip_with_parallel(y, 0.1, 0.9))
%timeit clip_with_parallel(y, 0.1, 0.9)
%timeit clip(y, 0.1, 0.9)
%timeit np.clip(y, 0.1, 0.9)
上述程序的運行結果將會是：
95.2 µs ± 5.6 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
104 µs ± 4.52 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
377 µs ± 62 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
這個慄子中的性能提升就是實打實的了。總之，使用parallel時不能一概而論，還是要做些實驗
需要指出的是，vectorize中的參數target一共有三種取值：cpu（默認）、parallel和cuda。關於選擇哪個取值，官方文檔上有很好的說明：
A general guideline is to choose different targets for different data sizes and algorithms. The 「cpu」 target works well for small data sizes (approx. less than 1KB) and low compute intensity algorithms. It has the least amount of overhead. The 「parallel」 target works well for medium data sizes (approx. less than 1MB). Threading adds a small delay. The 「cuda」 target works well for big data sizes (approx. greater than 1MB) and high compute intensity algorithms. Transfering memory to and from the GPU adds significant overhead.我們知道，Python 中由於有 GIL 的存在，所以多線程用起來非常不舒服。不過 numba 的 jit 裡面有一項參數叫 nogil，想來聰明的觀眾老爺們都猜到了它是幹什麼的了……
下面就來看一個慄子：
import math
from concurrent.futures import ThreadPoolExecutor
# 計算類似於 Sigmoid 的函數
def np_func(a, b):
return 1 / (a + np.exp(-b))
# 參數中的 result 代表的即是我們想要的結果，後同
# 第一個 kernel，nogil 參數設為了 False
@nb.jit(nopython=True, nogil=False)
def kernel1(result, a, b):
for i in range(len(result)):
result[i] = 1 / (a[i] + math.exp(-b[i]))
# 第二個 kernel，nogil 參數設為了 True
@nb.jit(nopython=True, nogil=True)
def kernel2(result, a, b):
for i in range(len(result)):
result[i] = 1 / (a[i] + math.exp(-b[i]))
def make_single_task(kernel):
def func(length, *args):
result = np.empty(length, dtype=np.float32)
kernel(result, *args)
return result
return func
def make_multi_task(kernel, n_thread):
def func(length, *args):
result = np.empty(length, dtype=np.float32)
args = (result,) + args
# 將每個線程接受的參數定義好
chunk_size = (length + n_thread - 1) // n_thread
chunks = [[arg[i*chunk_size:(i+1)*chunk_size] for i in range(n_thread)] for arg in args]
# 利用 ThreadPoolExecutor 進行並發
with ThreadPoolExecutor(max_workers=n_thread) as e:
for _ in e.map(kernel, *chunks):
pass
return result
return func
length = 10 ** 6
a = np.random.rand(length).astype(np.float32)
b = np.random.rand(length).astype(np.float32)
nb_func1 = make_single_task(kernel1)
nb_func2 = make_multi_task(kernel1, 4)
nb_func3 = make_single_task(kernel2)
nb_func4 = make_multi_task(kernel2, 4)
rs_np = np_func(a, b)
rs_nb1 = nb_func1(length, a, b)
rs_nb2 = nb_func2(length, a, b)
rs_nb3 = nb_func3(length, a, b)
rs_nb4 = nb_func4(length, a, b)
assert np.allclose(rs_np, rs_nb1, rs_nb2, rs_nb3, rs_nb4)
%timeit np_func(a, b)
%timeit nb_func1(length, a, b)
%timeit nb_func2(length, a, b)
%timeit nb_func3(length, a, b)
%timeit nb_func4(length, a, b)
這個慄子有點長，不過我們只需要知道如下兩點即可：
* make_single_task和make_multi_task分別生成單線程函數和多線程函數
* 生成的函數會調用相應的kernel來完成計算
上述程序的運行結果將會是：
14.9 ms ± 538 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.32 ms ± 259 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
10.2 ms ± 368 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
8.25 ms ± 279 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.68 ms ± 114 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
一般來說，數據量越大、並發的效果越明顯。反之，數據量小的時候，並發很有可能會降低性能
numba 的應用實例 —— 卷積與池化如果只想看效果的話倒沒什麼關係，不過如果想知道我具體在幹什麼的話，可以參見這篇文章
首先是卷積操作：
import numba as nb
import numpy as np
# 普通的卷積
def conv_kernel(x, w, rs, n, n_channels, height, width, n_filters, filter_height, filter_width, out_h, out_w):
for i in range(n):
for j in range(out_h):
for p in range(out_w):
window = x[i, ..., j:j+filter_height, p:p+filter_width]
for q in range(n_filters):
rs[i, q, j, p] += np.sum(w[q] * window)
return rs
# 簡單地加了個 jit 後的卷積
@nb.jit(nopython=True)
def jit_conv_kernel(x, w, rs, n, n_channels, height, width, n_filters, filter_height, filter_width, out_h, out_w):
for i in range(n):
for j in range(out_h):
for p in range(out_w):
window = x[i, ..., j:j+filter_height, p:p+filter_width]
for q in range(n_filters):
rs[i, q, j, p] += np.sum(w[q] * window)
# 卷積操作的封裝
def conv(x, w, kernel, args):
n, n_filters = args[0], args[4]
out_h, out_w = args[-2:]
rs = np.zeros([n, n_filters, out_h, out_w], dtype=np.float32)
kernel(x, w, rs, *args)
return rs
# 64 個 3 x 28 x 28 的圖像輸入（模擬 mnist）
x = np.random.randn(64, 3, 28, 28).astype(np.float32)
# 16 個 5 x 5 的 kernel
w = np.random.randn(16, x.shape[1], 5, 5).astype(np.float32)
n, n_channels, height, width = x.shape
n_filters, _, filter_height, filter_width = w.shape
out_h = height - filter_height + 1
out_w = width - filter_width + 1
args = (n, n_channels, height, width, n_filters, filter_height, filter_width, out_h, out_w)
print(np.linalg.norm((conv(x, w, conv_kernel, args) - conv(x, w, jit_conv_kernel, args)).ravel()))
%timeit conv(x, w, conv_kernel, args)
%timeit conv(x, w, jit_conv_kernel, args)
上述程序的運行結果將會是：
0.00112681
3.63 s ± 194 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
300 ms ± 20.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
可以看到，僅僅是加了一個jit、速度就直接提升了十多倍
有細心的觀眾老爺可能已經發現，我這裡沒有使用np.allclose；這是因為卷積涉及到的運算太多，僅僅是將數組的dtype從float64變成float32、精度就會大幅下降，所以使用np.allclose的話會過不了assert
同時需要特別注意的是，使用jit和使用純numpy進行編程的很大一點不同就是，不要畏懼用for；事實上一般來說，代碼「長得越像 C」、速度就會越快：
@nb.jit(nopython=True)
def jit_conv_kernel2(x, w, rs, n, n_channels, height, width, n_filters, filter_height, filter_width, out_h, out_w):
for i in range(n):
for j in range(out_h):
for p in range(out_w):
for q in range(n_filters):
for r in range(n_channels):
for s in range(filter_height):
for t in range(filter_width):
rs[i, q, j, p] += x[i, r, j+s, p+t] * w[q, r, s, t]
assert np.allclose(conv(x, w, jit_conv_kernel, args), conv(x, w, jit_conv_kernel, args))
%timeit conv(x, w, jit_conv_kernel, args)
%timeit conv(x, w, jit_conv_kernel2, args)
那麼程序的運行結果將會是：
281 ms ± 12.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
66.2 ms ± 2.32 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
可以看到這又快了 3 倍左右
接下來是池化操作（我們選用的是 MaxPool）：
# 普通的 MaxPool
def max_pool_kernel(x, rs, *args):
n, n_channels, pool_height, pool_width, out_h, out_w = args
for i in range(n):
for j in range(n_channels):
for p in range(out_h):
for q in range(out_w):
window = x[i, j, p:p+pool_height, q:q+pool_width]
rs[i, j, p, q] += np.max(window)
# 簡單地加了個 jit 後的 MaxPool
@nb.jit(nopython=True)
def jit_max_pool_kernel(x, rs, *args):
n, n_channels, pool_height, pool_width, out_h, out_w = args
for i in range(n):
for j in range(n_channels):
for p in range(out_h):
for q in range(out_w):
window = x[i, j, p:p+pool_height, q:q+pool_width]
rs[i, j, p, q] += np.max(window)
# 不懼用 for 的、「更像 C」的 MaxPool
@nb.jit(nopython=True)
def jit_max_pool_kernel2(x, rs, *args):
n, n_channels, pool_height, pool_width, out_h, out_w = args
for i in range(n):
for j in range(n_channels):
for p in range(out_h):
for q in range(out_w):
_max = x[i, j, p, q]
for r in range(pool_height):
for s in range(pool_width):
_tmp = x[i, j, p+r, q+s]
if _tmp > _max:
_max = _tmp
rs[i, j, p, q] += _max
# MaxPool 的封裝
def max_pool(x, kernel, args):
n, n_channels = args[:2]
out_h, out_w = args[-2:]
rs = np.zeros([n, n_filters, out_h, out_w], dtype=np.float32)
kernel(x, rs, *args)
return rs
pool_height, pool_width = 2, 2
n, n_channels, height, width = x.shape
out_h = height - pool_height + 1
out_w = width - pool_width + 1
args = (n, n_channels, pool_height, pool_width, out_h, out_w)
assert np.allclose(max_pool(x, max_pool_kernel, args), max_pool(x, jit_max_pool_kernel, args))
assert np.allclose(max_pool(x, jit_max_pool_kernel, args), max_pool(x, jit_max_pool_kernel2, args))
%timeit max_pool(x, max_pool_kernel, args)
%timeit max_pool(x, jit_max_pool_kernel, args)
%timeit max_pool(x, jit_max_pool_kernel2, args)
上述程序的運行結果將會是：
586 ms ± 38 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
8.25 ms ± 526 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.4 ms ± 57 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
可以看到最快的比最慢的要快整整 400 倍有多
以上就是 numba 的基本應用。