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

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

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

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

Numpy ”愛莫能助“的時刻

假設你想要將一個非常大的數組轉變為按遞增順序排序：很好理解，就是將元素按值的大小升序排列，如：

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

以下是一個簡單的就地轉換方式：

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 運行很快，是因為它可以在不調用 python 自身解釋器的前提下完成所有計算。但對于上面這個場景(python 中的循環)，就會暴露出一個問題：我們會失去 Numpy 得天獨厚的性能優勢。

對一個含有一千萬個元素的 Numpy 數組使用上面的函數進行轉換，在我的電腦上需要運行 2.5 秒。那么，還可以優化得更快嗎？

使用 Numba 提速

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

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

再次運行，發現僅需要 0.19 秒，在完全重用舊代碼邏輯的前提下，感覺效果還不錯。

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

Numba 簡介

在 Numpy 或 Scipy 中找到目標函數，可以很快解決常見的計算問題。但是如果函數不存在呢?(比如剛剛的 numpy.maximum.accumulate)。這種情況下如果想加速代碼運行。可能會選擇其他低級的編程語言來實現擴展，但這也意味著切換編程語言，會讓模塊構建和系統總體變得更復雜。

使用 Numba 你可以做到：

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

Numba 首先會解析代碼，然后根據數據的輸入類型以即時的方式編譯它們。例如，當輸入是 u64 數組和浮點型數組時，分別得到的編譯結果是不一樣的。

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

Numba 的一些短板

需要一次代碼編譯耗時

當第一次調用 Numba 修飾的函數時，它需要花費一定的時間來生成對應的機器代碼。比如，我們可以使用 IPython 的 %time 命令來計算運行一個 Numba 修飾的函數需要花費多長時間：

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

可以看到，函數第一次調用后運行非常慢(注意單位時毫秒而不是微秒)，這就是因為它需要時間來編譯生成機器代碼。不過函數后面的運行速度會顯著提升。這種時間成本在輸入數據的類型發生變化時會再次消耗，比如，我們將輸入類型換為浮點數：

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

計算兩數之和當然不需要啟用 Numba，這里用這個案例是因為能夠比較容易地看出編譯所需的時間成本。

與 python 和 Numpy 的不同實現方式

Numba 在功能方面可以說是實現了 python 的一個子集，也可以說是實現了 Numpy API 的一個子集，這將會導致一些潛在的問題：

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

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

• 由于使用的不用的算法，兩者的性能表現會有區別
• 可能會由于 bug 導致結果不一致

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

Numba 與其他選項的對比

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

結語

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