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

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

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

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

Numpy ”愛(ài)莫能助“的時(shí)刻

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

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

以下是一個(gè)簡(jiǎn)單的就地轉(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 自身解釋器的前提下完成所有計(jì)算。但對(duì)于上面這個(gè)場(chǎng)景(python 中的循環(huán))，就會(huì)暴露出一個(gè)問(wèn)題：我們會(huì)失去 Numpy 得天獨(dú)厚的性能優(yōu)勢(shì)。

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

使用 Numba 提速

Numba 是一款為 python 打造的、專門(mén)針對(duì) Numpy 數(shù)組循環(huán)計(jì)算場(chǎng)景的即時(shí)編譯器。顯然，這正是我們所需要的。讓我們?cè)谠泻瘮?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 秒，在完全重用舊代碼邏輯的前提下，感覺(jué)效果還不錯(cuò)。

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

Numba 簡(jiǎn)介

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

使用 Numba 你可以做到：

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

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

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

Numba 的一些短板

需要一次代碼編譯耗時(shí)

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

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)行非常慢(注意單位時(shí)毫秒而不是微秒)，這就是因?yàn)樗枰獣r(shí)間來(lái)編譯生成機(jī)器代碼。不過(guò)函數(shù)后面的運(yùn)行速度會(huì)顯著提升。這種時(shí)間成本在輸入數(shù)據(jù)的類型發(fā)生變化時(shí)會(huì)再次消耗，比如，我們將輸入類型換為浮點(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

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

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

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

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

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

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

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

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

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

結(jié)語(yǔ)

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