【实用小功能10】python运行加速神器——numba（详细教学版）

1. 为什么python这么慢？

python比c++慢，尤其是存在循环的情况下，python和c+的区别主要有：

1.1 动态变量

c++中需要对变量类型有严格的定义，比如int或者float类型。但是python不需要，它去掉了变量申明和数据类型。python是一种动态类型语言，它会在做运算的时候根据变量的值自动推断出其类型，这种动态类型的特性使得python编写更加简单和灵活，但同时增加了一定的时间开销，反观c/c++在运算时则只是简单的内存读写和机器指令（加减乘除…），所以c/c++的速度比python快很多。

1.2. 解释性语言

c/c++是编译型语言，需要先将源代码编译成机器码，然后再运行程序。

python是解释性语言，它的源代码是直接由解释器解释执行的，这意味着在编写python代码时，不需要先将代码编译成可执行文件，而是可以直接运行代码。所以python的执行速度相较于c/c++会慢很多。

这个问题的一种解决方法就是即时编译JIT（Just-in-time compilation）：JIT编译器会动态地将高级语言编写的代码转换为机器码，可以直接由计算机的处理器执行，这是在运行时完成的，也就是代码执行之前，因此称为“即时”。JIT针对特定的硬件和操作系统进行代码优化，可以使得python代码获得显著的性能提升。

2. Numba的介绍和使用

github主页：https://github.com/numba/numba

在线文档：https://numba.readthedocs.io/en/stable/index.html

可直接pip/conda下载库

conda install numba
pip install numba

numba非常适合于使用了numpy数组、函数和循环的代码，使用的方法就是装饰器，用它！用的好时间节省个几十倍。numba的可使用范围：

• 操作系统：Windows (64 bit), OSX, Linux (64 bit).
• 架构：x86, x86_64, ppc64le, armv8l (aarch64), M1/Arm64.
• GPU：Nvidia CUDA.
• CPython
• NumPy版本：1.22~1.25

2.1 numba加速python小实例

• 使用python计算一个矩阵的所有元素的和

• 原始代码

def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result 

start = time.perf_counter()
a = np.random.random((5000, 5000)) 
result = cal_sum(a)
end = time.perf_counter()
print("原始代码耗时：{}s".format((end - start)))

#OUT:
#原始代码耗时：5.725140199996531s

• 使用numba加速

import numba
 
@numba.jit(nopython=True) 
def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result 

start = time.perf_counter()
a = np.random.random((5000, 5000)) 
result = cal_sum(a)
end = time.perf_counter()
print("加速后耗时：{}s".format((end - start)))

#OUT:
#加速后耗时：0.2892118000017945s

@jit(nopython=True)与@njit是等同的，这个选项指示numba在编译时尽可能地避免使用python对象，将python代码转换为更快的机器码，而不是使用python解释器来执行代码。

• 多次调用同一个使用了numba加速的函数

numba在第首次调用函数的时候进行了编译，但当编译发生后，numba会将该函数的机器码进行缓存，如果再次调用该函数，它会直接从缓存中加载，而不需要再次编译。简而言之，在一段代码中的使用了numba的函数，第二次、第三次…调用该函数运算都要比第一次耗时更少。

x = np.arange(100).reshape(10, 10)

@jit(nopython=True)
def go_fast(a):
    trace = 0.0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("首次调用耗时：{}s".format((end - start)))

start = time.perf_counter()
go_fast(x)
end = time.perf_counter()
print("再次调用耗时：{}s".format((end - start)))

#OUT:
#首次调用耗时：0.16958359999989625s
#再次调用耗时：4.7599998652003706e-05s

• 使用并行加速【失败，原因还未知】

官方使用说明网址：numba-parallel

拿了官方的代码对比了下未使用并行和使用并行的，但是并行的反而更慢，没仔细细究缘由。

from numba import njit

@njit()
def prange_test(A):
    s = 0
    for i in range(A.shape[0]):
        s += A[i]
    return s

start = time.perf_counter()
a = np.arange(1000)
result = prange_test(a)
end = time.perf_counter()
print("无并行耗时：{}s".format((end - start)))

#OUT:
#无并行耗时：0.07213479999336414s

from numba import njit, prange

@njit(parallel=True)
def prange_test(A):
    s = 0
    # Without "parallel=True" in the jit-decorator
    # the prange statement is equivalent to range
    for i in prange(A.shape[0]):
        s += A[i]
    return s

start = time.perf_counter()
a = np.arange(1000)
result = prange_test(a)
end = time.perf_counter()
print("并行耗时：{}s".format((end - start)))

#OUT:
#并行耗时：0.36391139999614097s

2.2 个人经验

上面的函数比较简单，但是在实际使用中我们的函数可能复杂很多，简单地直接在函数上添加一个装饰器会出现各种报错，我们需要针对问题调整函数代码，大概陈列几个我遇到的问题及解决的方法。

• 报错1

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Cannot infer the type of variable 'result', have imprecise type: list(undefined)<iv=None>. 

For Numba to be able to compile a list, the list must have a known and
precise type that can be inferred from the other variables. Whilst sometimes
the type of empty lists can be inferred, this is not always the case, see this
documentation for help:

https://numba.pydata.org/numba-doc/latest/user/troubleshoot.html#my-code-has-an-untyped-list-problem

有些报错，会给个网址告诉你报错原因，点进去认真看下然后照着意思修改修改自己的代码就行了，但这个问题的解决方法如下：

@numba.jit(nopython=True) 
def test(): 
    result, result1 = [], 0
    return result, result1

test()

# 改为
@numba.jit(nopython=True) 
def test(): 
    result = []
    result1 = 0
    return result, result1

test()

• 报错2：

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Unknown attribute 'tolist' of type array(float32, 1d, C)
 ...

Untyped global name 'input': Cannot determine Numba type of <class 'method'>

这样的问题是因为numba现在还有一些函数是不支持，比如上面python自带的input()函数和数组的tolist()函数，遇到这些问题的时候解决方法就是使用同样功能的且numba支持的函数or自己重写一个函数。

• 报错3

def test_func():
    return True

@jit(nopython=True)
def test(): 
    flag = test_func()

test()

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'test_func': Cannot determine Numba type of <class 'function'>

这个问题是numba不支持函数中嵌套自己的函数，解决方法就是把调用的函数合并到一个函数中。

• 报错4

会有一些与list列表相关的报错，numba貌似很不支持列表操作，尤其是嵌套列表类似[[[1,2],[2,3]],[[2,2],[8,3]]]，解决方法是将这种列表改为数组进行操作。

2.3 其他

numba还提供了很多其他加速方法（我还没试）：

• 想节省编译时间，numba提供了提前编译模式：Ahead-of-Time compilation
• 使用GPU：cuda-index

参考：Python 提速大杀器之 numba 篇