python sum函數numpy_如何用numba加速python？

我把寫好的markdown導入進來，但是沒想到知乎的排版如此感人。如果對知乎排版不滿想要看高清清爽版，請移步微信公眾號原文如何用numba加速python？同時歡迎關注

前言

說道現在最流行的語言，就不得不提python。可是python雖然容易上手，但速度卻有點感人。如何用簡單的方法讓python加速到近乎可以媲美C的速度呢？今天來就來談談numba這個寶貝。對你沒看錯，不是numpy，就是numba。

用函數編程

在面對一個計算project的時候，我們最容易想到的就是直接碼代碼，最后寫出一個超長的程序。這樣一來，一旦出錯往往就需要花很多時間定位問題。

有一個簡單的辦法解決這個問題，就是定義各種各樣的函數，把任務分解成很多小部分。因為每個函數都不是特別復雜，并且在寫好的時候就可以隨時檢查，因此簡潔的主程序一旦出問題就很容易定位并解決。面向對象編程的思想就是基于函數。

寫好函數之后，還可以使用裝飾器（decorator）讓它變得強大。裝飾器本身是一個函數，不過是函數的函數，目的是增加函數的功能。比如首先定義一個輸出當前時間的函數，再定義一個規定時間格式的函數，把后一個函數作用在前一個函數上，就是一個裝飾器，作用是用特定格式輸出當前時間。

Numba的優勢

1.簡單，往往只要1行代碼就有驚喜；

2.對循環（loop）有奇效，而往往在科學計算中限制python速度的就是loop；

3.兼容常用的科學計算包，如numpy、cmath等；

4.可以創建ufunc；

5.會自動調整精度，保證準確性。

如何使用numba

針對上面提到的numba的優勢，我來進行逐一介紹。首先導入numba

import numba as nb

只用1行代碼即可加速，對loop有奇效

因為numba內置的函數本身是個裝飾器，所以只要在自己定義好的函數前面加個@nb.jit()就行，簡單上手。下面以一個求和函數為例

# 用numba加速的求和函數
@nb.jit()
def nb_sum(a):Sum = 0for i in range(len(a)):Sum += a[i]return Sum# 沒用numba加速的求和函數
def py_sum(a):Sum = 0for i in range(len(a)):Sum += a[i]return Sum

來測試一下速度

import numpy as np
a = np.linspace(0,100,100) # 創建一個長度為100的數組%timeit np.sum(a) # numpy自帶的求和函數
%timeit sum(a) # python自帶的求和函數
%timeit nb_sum(a) # numba加速的求和函數
%timeit py_sum(a) # 沒加速的求和函數

結果如下

# np.sum(a)
7.1 μs ± 537 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
# sum(a)
27.7 μs ± 2.64 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
# nb_sum(a)
1.05 μs ± 27.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
# py_sum(a)
43.7 μs ± 1.71 μs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

可以看出，numba甚至比號稱最接近C語言速度運行的numpy還要快6倍以上。但大家都知道，numpy往往對大的數組更加友好，那我們來測試一個更長的數組

a = np.linspace(0,100,10**6) # 創建一個長度為100萬的數組

測試結果如下

# np.sum(a)
2.51 ms ± 246 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# sum(a)
249 ms ± 19.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
# nb_sum(a)
3.01 ms ± 59.7 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# py_sum(a)
592 ms ± 42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

可見即便是用很長的loop來計算，numba的表現也絲毫不亞于numpy。在這里，我們可以看到numba相對于numpy一個非常明顯的優勢：numba可以把各種具有很大loop的函數加到很快的速度，但numpy的加速只適用于numpy自帶的函數。

但要注意的是，numba對沒有循環或者只有非常小循環的函數加速效果并不明顯，用不用都一樣。（偷偷告訴你，numba的loop甚至常常比numpy的矩陣運算還要快）

兼容常用的科學計算包，可以創建ufunc

上一部分我們比較了numba和numpy的表現，可以說numba非常亮眼了。但numpy還有一個非常強大的功能——ufunc (universal functions)，它可以讓一個函數同時處理很多數據。比如要求一個數組每一個元素的三角函數，只需要

np.sin(a) # 這里的a仍然是上面有100萬個元素的數組

而不需要寫個循環一個一個求。可如果不用numpy但又想要很快的速度，那應該怎么求呢？我們可以用從math庫里導入sin，然后寫個loop再用numba加速。除了這個方法，在這里我還想說numba另一個強大的功能，矢量化（vectorize）。像上面的jit一樣，只要添加一行vectorize就可以讓普通函數變成ufunc

from math import sin@nb.vectorize()
def nb_vec_sin(a):return sin(a)

來比較一下用各種方式寫出的三角函數

# 用numba加速的loop
13.5 ms ± 405 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# nb_vec_sin(a)
14.2 ms ± 55.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# np.sin(a)
5.75 ms ± 85 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

可以看出，用vectorize改寫的三角函數具有了和np.sin()一樣的同時處理數組每個元素的功能，而且表現也不必numpy差。當遇到numpy沒有的數學函數時（比如sech），用numba矢量化不失為一個好的選擇。除了math中的sin，它支持的其他函數列表可以在documentation中找到（鏈接見附錄）。

其實numpy也有矢量化的功能，只不過比numba差遠了。

會自動調整精度，保證準確性

上面我們用的測試數組數字范圍只是從0到100的。可如果數字很大，那么就很容易出現overflow的問題，比如

a = np.arange(10**6) # a的最小值為0，最大值為10**6-1

你猜猜用python自帶的sum，我們自己寫的py_sum，np.sum和nb_sum給出的結果一不一樣呢？你會發現

# np.sum(a)
1783293664
# sum(a)
1783293664
# nb_sum(a)
499999500000
# py_sum(a)
1783293664

numba的結果和其他三個都不一樣，肯定錯了呀，還用問么？

且慢！其實在運行的時候，我并沒有告訴你sum和py_sum都報錯了“RuntimeWarning: overflow encountered in long_scalars”。但奇怪的是np.sum并沒有報錯。

在上面的四個函數里，其實numba表現的最好，因為它自動調整了整數類型。如果你用nb.typeof()查看，你會發現numba給出的結果是int64，而其他三個都是int32。不得不說，numba不僅快還在精度方面表現很好！

拓展

在本文的最后一部分，我想談兩個問題。

Numba的精度問題

精度方面，在上面我也談到numba會自動轉換數據類型以適應計算。但是在個別時候，這種自動轉變類型可能會引起一些計算誤差。通常這個誤差是非常小的，幾乎不會造成任何影響。但如果你所處理的問題會積累誤差，比如求解非線性方程，那么在非常多的計算之后誤差可能就是肉眼可見了。如果你發現有這樣的問題，記得在jit中指定輸入輸出的數據類型。numba具有C所有的數據類型，比如對上面的求和函數，只需要把@nb.jit()改為@nb.jit(nb.int64(nb.int32[:]))即可。nb.int64是說輸出的數字為int64類型，nb.int32是說輸入的數據類型為int32，而[:]是說輸入的是數組。