Python 本身是一門運(yùn)行較慢的語言，因此對于計算場景，最好的優(yōu)化方式就是優(yōu)化代碼寫法。你可以使用現(xiàn)有的科學(xué)計算庫：比如 Numpy 和 Scipy。但如果想要在不使用低級語言(如 CPython、Rust 等)實(shí)現(xiàn)擴(kuò)展的前提下實(shí)現(xiàn)一個新的算法時，該如何做呢?

對于某些特定的、尤其是針對數(shù)組的計算場景，Numba 可以顯著加快代碼的運(yùn)行速度。在使用時，我們有時候需要調(diào)整一下原始代碼，而有時候卻又不需要做任何改動。當(dāng)它真正起到作用時，效果將會非常明顯。

在本篇文章中，我們會談及以下幾方面：

• 為什么 有時候單獨(dú)使用 Numpy 是不夠的
• Numba 的基礎(chǔ)使用方式
• Numba 是如何在很高的層次上來對你的代碼運(yùn)行造成影響的

Numpy ”愛莫能助“的時刻

假設(shè)你想要將一個非常大的數(shù)組轉(zhuǎn)變?yōu)榘催f增順序排序：很好理解，就是將元素按值的大小升序排列，如：

[1, 2, 1, 3, 3, 5, 4, 6] → [1, 2, 2, 3, 3, 5, 5, 6]

以下是一個簡單的就地轉(zhuǎn)換方式：

def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value

Numpy 運(yùn)行很快，是因?yàn)樗梢栽诓徽{(diào)用 python 自身解釋器的前提下完成所有計算。但對于上面這個場景(python 中的循環(huán))，就會暴露出一個問題：我們會失去 Numpy 得天獨(dú)厚的性能優(yōu)勢。

對一個含有一千萬個元素的 Numpy 數(shù)組使用上面的函數(shù)進(jìn)行轉(zhuǎn)換，在我的電腦上需要運(yùn)行 2.5 秒。那么，還可以優(yōu)化得更快嗎？

使用 Numba 提速

Numba 是一款為 python 打造的、專門針對 Numpy 數(shù)組循環(huán)計算場景的即時編譯器。顯然，這正是我們所需要的。讓我們在原有函數(shù)的基礎(chǔ)上添加兩行代碼試試：

from numba import njit

@njit
def monotonically_increasing(a):
    max_value = 0
    for i in range(len(a)):
        if a[i] > max_value:
            max_value = a[i]
        a[i] = max_value

再次運(yùn)行，發(fā)現(xiàn)僅需要 0.19 秒，在完全重用舊代碼邏輯的前提下，感覺效果還不錯。

實(shí)際上 Numpy 也有一個特殊的函數(shù)可以解決這種場景(但是會修改原有函數(shù)的代碼邏輯)：`numpy.maximum.accumulate` 。通過使用它，函數(shù)的運(yùn)行時長會縮短至 0.03 秒。

Numba 簡介

在 Numpy 或 Scipy 中找到目標(biāo)函數(shù)，可以很快解決常見的計算問題。但是如果函數(shù)不存在呢?(比如剛剛的 numpy.maximum.accumulate)。這種情況下如果想加速代碼運(yùn)行?？赡軙x擇其他低級的編程語言來實(shí)現(xiàn)擴(kuò)展，但這也意味著切換編程語言，會讓模塊構(gòu)建和系統(tǒng)總體變得更復(fù)雜。

使用 Numba 你可以做到：

• 使用 python 和擁有更快編譯速度的解釋器運(yùn)行同一份代碼
• 簡單快速地迭代算法

Numba 首先會解析代碼，然后根據(jù)數(shù)據(jù)的輸入類型以即時的方式編譯它們。例如，當(dāng)輸入是 u64 數(shù)組和浮點(diǎn)型數(shù)組時，分別得到的編譯結(jié)果是不一樣的。

Numba 還可以對非 CPU 的計算場景生效：比如你可以 在 GPU 上運(yùn)行代碼。誠然，上文中的示例只是 Numba 的一個最小應(yīng)用，官方文檔中還有很多特性可供選擇。

Numba 的一些短板

需要一次代碼編譯耗時

當(dāng)?shù)谝淮握{(diào)用 Numba 修飾的函數(shù)時，它需要花費(fèi)一定的時間來生成對應(yīng)的機(jī)器代碼。比如，我們可以使用 IPython 的 %time 命令來計算運(yùn)行一個 Numba 修飾的函數(shù)需要花費(fèi)多長時間：

In [1]: from numba import njit

In [2]: @njit
   ...: def add(a, b): a + b

In [3]: %time add(1, 2)
CPU times: user 320 ms, sys: 117 ms, total: 437 ms
Wall time: 207 ms

In [4]: %time add(1, 2)
CPU times: user 17 μs, sys: 0 ns, total: 17 μs
Wall time: 24.3 μs

In [5]: %time add(1, 2)
CPU times: user 8 μs, sys: 2 μs, total: 10 μs
Wall time: 13.6 μs

可以看到，函數(shù)第一次調(diào)用后運(yùn)行非常慢(注意單位時毫秒而不是微秒)，這就是因?yàn)樗枰獣r間來編譯生成機(jī)器代碼。不過函數(shù)后面的運(yùn)行速度會顯著提升。這種時間成本在輸入數(shù)據(jù)的類型發(fā)生變化時會再次消耗，比如，我們將輸入類型換為浮點(diǎn)數(shù)：

In [8]: %time add(1.5, 2.5)
CPU times: user 40.3 ms, sys: 1.14 ms, total: 41.5 ms
Wall time: 41 ms

In [9]: %time add(1.5, 2.5)
CPU times: user 16 μs, sys: 3 μs, total: 19 μs
Wall time: 26 μs

計算兩數(shù)之和當(dāng)然不需要啟用 Numba，這里用這個案例是因?yàn)槟軌虮容^容易地看出編譯所需的時間成本。

與 python 和 Numpy 的不同實(shí)現(xiàn)方式

Numba 在功能方面可以說是實(shí)現(xiàn)了 python 的一個子集，也可以說是實(shí)現(xiàn)了 Numpy API 的一個子集，這將會導(dǎo)致一些潛在的問題：

(1)會出現(xiàn) python 和 Numpy 部分特性都不支持的情況

(2)由于 Numba 重新實(shí)現(xiàn)了 Numpy 的 API，在使用時可能會出現(xiàn)以下情況：

• 由于使用的不用的算法，兩者的性能表現(xiàn)會有區(qū)別
• 可能會由于 bug 導(dǎo)致結(jié)果不一致

(3)另外，當(dāng) Numba 編譯失敗時，其暴露的錯誤信息可能會很難理解

Numba 與其他選項(xiàng)的對比

• 僅使用 Numpy 和 Scipy：可以讓 python 代碼運(yùn)行時達(dá)到其他語言編譯器的速度，但是對于某些循環(huán)計算的場景不生效
• 直接使用低級語言編寫代碼：這意味著你可以優(yōu)化所有的代碼語句，但是需要拋棄 python 使用另一門語言
• 使用 Numba：可以優(yōu)化 python 循環(huán)計算的場景，但是對于某些 python 語言本身和 Numpy API 的特性使用會受到限制

結(jié)語

Numba 最棒的地方在于嘗試起來非常簡單。因此每當(dāng)你有一個做一些數(shù)學(xué)運(yùn)算且運(yùn)行緩慢的 for 循環(huán)時，可以嘗試使用 Numba ：運(yùn)氣好的話，它只需要兩行代碼就可以顯著加快代碼運(yùn)行速度。