Python程序提速神器--Numba快速上手指南

程序员文章站 2022-06-17 16:55:45

...

前言

如果你在使用 Python 进行高性能计算， Numba 提供的加速效果可以比肩原生的 C/C++ 程序，只需要在函数上添加一行 @jit 的装饰。它支持 CPU 和 GPU ，是数据科学家必不可少的编程利器。

我们知道，计算机只能执行二进制的机器码，C、C++ 等编译型语言依靠编译器将源代码转化为可执行文件后才能运行，Python、Java 等解释型语言使用解释器将源代码翻译后在虚拟机上执行。对于 Python ，由于解释器的存在，其执行效率比 C 语言慢几倍甚至几十倍。
Python程序提速神器--Numba快速上手指南
上图比较了当前流行的各大编程语言在几个不同任务上的计算速度。C 语言经过几十年的发展，优化已经达到了极致。以 C 语言为基准，大多数解释语言，如 Python、R 会慢十倍甚至一百倍。 Julia 这个解释语言是个“奇葩”，因为它采用了 JIT 编译技术。

解决 Python 执行效率低的问题，一种解决办法是使用 C/C++ 语言重写 Python 函数，但是这要求程序员对 C/C++ 语言熟悉，且调试速度慢，不适合绝大多数 Python 程序员。另外一种非常方便快捷的解决办法就是使用 Just-In-Time（JIT） 技术，本文将解释 JIT 技术的原理，并提供几个案例，让你十分钟内学会 JIT 技术。

一、Python解释器的工作原理

Python 是一门解释语言， Python 为我们提供了基于硬件和操作系统的一个虚拟机，并使用解释器将源代码转化为虚拟机可执行的字节码。字节码在虚拟机上执行，得到结果。
Python程序提速神器--Numba快速上手指南
我们使用 python example.py 来执行一份源代码时， Python 解释器会在后台启动一个字节码编译器（ Bytecode Compiler ），将源代码转换为字节码。字节码是一种只能运行在虚拟机上的文件， Python 的字节码默认后缀为 .pyc ， Python 生成 .pyc 后一般放在内存中继续使用，并不是每次都将 .pyc 文件保存到磁盘上。有时候我们会看到自己 Python 代码文件夹里有很多 .pyc 文件与 .py 文件同名，但也有很多时候看不到 .pyc 文件。pyc 字节码通过 Python 虚拟机与硬件交互，关于.pyc文件的作用和更新机制请移步Python程序的执行过程（解释型语言和编译型语言）。虚拟机的出现导致程序和硬件之间增加了中间层，运行效率大打折扣。相信使用过虚拟机软件的朋友深有体会，在原生的系统上安装一个虚拟机软件，在虚拟机上再运行一个其他系统，经常感觉速度下降，体验变差，这与 Python 虚拟机导致程序运行慢是一个原理。

Just-In-Time（JIT） 技术为解释语言提供了一种优化，它能克服上述效率问题，极大提升代码执行速度，同时保留 Python 语言的易用性。使用 JIT 技术时，JIT 编译器将 Python 源代码编译成机器直接可以执行的机器语言，并可以直接在 CPU 等硬件上运行。这样就跳过了原来的虚拟机，执行速度几乎与用 C 语言编程速度并无二致。

二、快速上手Numba

Numba 是一个针对 Python 的开源 JIT 编译器，由 Anaconda 公司主导开发，可以对原生代码进行 CPU 和 GPU 加速。

使用 conda 安装 Numba ：

conda install numba

或者使用 pip 安装：

pip install numba

使用时，只需要在原来的函数上添加一行"注释"：

from numba import jit
import numpy as np

SIZE = 2000
x = np.random.random((SIZE, SIZE))

"""
给定n*n矩阵，对矩阵每个元素计算tanh值，然后求和。
因为要循环矩阵中的每个元素，计算复杂度为 n*n。
"""
@jit
def jit_tan_sum(a):   # 函数在被调用时编译成机器语言
    tan_sum = 0
    for i in range(SIZE):   # Numba 支持循环
        for j in range(SIZE):
            tan_sum += np.tanh(a[i, j])   # Numba 支持绝大多数NumPy函数
    return tan_sum

%timeit print(jit_tan_sum(x))  # 添加numba的jit装饰器之后的测试（1）  

def jit_tan_sum_normal(a):
    tan_sum = 0
    for i in range(SIZE):
        for j in range(SIZE):
            tan_sum += np.tanh(a[i, j])
    return tan_sum

%timeit print(jit_tan_sum_normal(x))  # 未添加numba的jit装饰器之后的测试（2）

我们只需要在原来的代码上添加一行 @jit ，即可将一个函数编译成机器码，其他地方都不需要更改。 @ 符号装饰了原来的代码，所以称类似写法为装饰器。

在我的 Core i5 处理器上，使用 jupyter notebook 编写测试代码，添加 @jit 装饰器前后的运行结果为：

测试①：
Python程序提速神器--Numba快速上手指南
这是未添加 jit 装饰器时的结果。

测试②：
Python程序提速神器--Numba快速上手指南
这是添加了 jit 装饰器之后的结果。

可以看到这里使用 Numba 加速后代码的执行速度提升了 169 倍！而且随着数据和计算量的增大，numba 的性能提升可能会更大！

PS：不同的机器运行结果也会有所差异，此运行结果最终解释权仅归本人所有！！！

三、Numba的使用场景

Numba 简单到只需要在函数上加一个装饰器就能加速程序，但也有缺点。目前 Numba 只支持了 Python 原生函数和部分 NumPy 函数，其他一些场景可能不适用。比如 Numba 官方给出这样的例子：

from numba import jit
import pandas as pd

x = {'a': [1, 2, 3], 'b': [20, 30, 40]}

@jit
def use_pandas(a): # Function will not benefit from Numba jit
    df = pd.DataFrame.from_dict(a) # Numba doesn't know about pd.DataFrame
    df += 1                        # Numba doesn't understand what this is
    return df.cov()                # or this!

%timeit print(use_pandas(x))  # 测试（3）  

def use_pandas_normal(a): 
df = pd.DataFrame.from_dict(a) 
df += 1                        
return df.cov()                

%timeit print(use_pandas_normal(x))  # 测试（4）

测试③的结果为：
Python程序提速神器--Numba快速上手指南

测试④的结果为：
Python程序提速神器--Numba快速上手指南
可以看见两次测试的结果相差无几，为什么呢？

pandas 是更高层次的封装， Numba 其实不能理解它里面做了什么，所以无法对其加速。一些大家经常用的机器学习框架，如 scikit-learn ， tensorflow ， pytorch 等，已经做了大量的优化，不适合再使用 Numba 做加速。

此外，Numba 还不支持：

try ... except 异常处理；
with 语句；
yield from 语句。

那么 Numba 当前支持的功能有哪些呢？
详情请移步：http://numba.pydata.org/numba-doc/latest/reference/pysupported.html

那如何决定是否使用 Numba 呢？

Numba 的 @jit 装饰器就像自动驾驶，用户不需要关注到底是如何优化的， Numba 去尝试进行优化，如果发现不支持，那么 Numba 会继续用 Python 原来的方法去执行该函数，即图 Python解释器工作原理 中左侧部分。这种模式被称为 object 模式。前文提到的 pandas 的例子， Numba 发现无法理解里面的内容，于是自动进入了 object 模式。object 模式还是和原生的 Python 一样慢，还有可能比原来更慢。

Numba 真正牛逼之处在于其 nopython 模式。将装饰器改为 @jit(nopython=True) 或者 @njit ， Numba 会假设你已经对所加速的函数非常了解，强制使用加速的方式，不会进入 object 模式，如编译不成功，则直接抛出异常。 nopython 的名字会有点歧义，我们可以理解为不使用很慢的 Python ，强制进入图 Python解释器工作原理 中右侧部分。

实践上，一般推荐将代码中计算密集的部分作为单独的函数提出来，并使用 nopython 方式优化，这样可以保证我们能使用到 Numba 的加速功能。其余部分还是使用 Python 原生代码，在计算加速的前提下，避免过长的编译时间。（有关编译时间的问题下节将会介绍。） Numba 可以与 NumPy 紧密结合，两者一起，常常能够得到近乎 C 语言的速度。尽管 Numba 不能直接优化 pandas ，但是我们可以将 pandas 中处理数据的 for 循环作为单独的函数提出来，再使用 Numba 加速。

四、Numba的编译开销

编译源代码需要一定的时间。C/C++ 等编译型语言要提前把整个程序先编译好，再执行可执行文件。 Numba 库提供的是一种 惰性编译（Lazy Compilation） 技术，即在运行过程中第一次发现代码中有 @jit ，才将该代码块编译。用到的时候才编译，看起来比较懒，所以叫惰性编译。使用 Numba 时， 总时间 = 编译时间 + 运行时间 。相比所能节省的计算时间，编译的时间开销很小，所以物有所值。对于一个需要多次调用的 Numba 函数，只需要编译一次，后面再调用时就不需要编译了。

from numba import jit
import numpy as np
import time

SIZE = 2000
x = np.random.random((SIZE, SIZE))

"""
给定n*n矩阵，对矩阵每个元素计算tanh值，然后求和。
因为要循环矩阵中的每个元素，计算复杂度为 n*n。
"""
@jit
def jit_tan_sum(a):   # 函数在被调用时编译成机器语言
    tan_sum = 0
    for i in range(SIZE):   # Numba 支持循环
        for j in range(SIZE):
            tan_sum += np.tanh(a[i, j])   # Numba 支持绝大多数NumPy函数
    return tan_sum

# 总时间 = 编译时间 + 运行时间
start = time.time()
jit_tan_sum(x)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

# Numba将加速的代码缓存下来
# 总时间 = 运行时间
start = time.time()
jit_tan_sum(x)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

运行结果为：
Python程序提速神器--Numba快速上手指南
代码中两次调用 Numba 优化函数，第一次执行时需要编译，第二次使用缓存的代码，运行时间将大大缩短。

原生 Python 速度慢的另一个重要原因是变量类型不确定。声明一个变量的语法很简单，如 a = 1 ，但没有指定 a 到底是一个整数和一个浮点小数。 Python 解释器要进行大量的类型推断，会非常耗时。同样，引入 Numba 后， Numba 也要推断输入输出的类型，才能转化为机器码。针对这个问题， Numba 给出了名为 Eager Compilation 的优化方式。

from numba import jit, int32

@jit("int32(int32, int32)", nopython=True)
def f2(x, y):
    # A somewhat trivial example
    return x + y  

%timeit f2(1000, 100000)  # 测试（5）

@jit
def f2(x, y):
    # A somewhat trivial example
    return x + y

%timeit f2(1000, 100000)  # 测试（6）

测试⑤和⑥的结果分别为：
Python程序提速神器--Numba快速上手指南

@jit(int32(int32, int32)) 告知 Numba 你的函数在使用什么样的输入和输出，括号内是输入，括号左侧是输出。这样不会加快执行速度，但是会加快编译速度，可以更快将函数编译到机器码上。

五、Numba到底有多快

网上有很多对 Numba 进行性能评测的文章，在一些计算任务上， Numba 结合 NumPy ，可得到接近 C 语言的速度。
Python程序提速神器--Numba快速上手指南

六、Numba的更多功能

除了上面介绍的加速功能，Numba 还有很多其他功能。 @vectorize 装饰器可以将一个函数向量化，变成类似 NumPy 函数一样，直接处理矩阵和张量。R 语言用户可能非常喜欢这个功能。

Numba 还可以使用 GPU 进行加速，目前支持英伟达的 CUDA 和 AMD 的 ROC 。 GPU 的工作原理和编程方法与 CPU 略有不同。

七、Numba原理

Python程序提速神器--Numba快速上手指南
Numba 使用了 LLVM 和 NVVM 技术，这个技术可以将 Python、Julia 这样的解释语言直接翻译成 CPU 或 GPU 可执行的机器码。

总结

无论你是在做 NLP 自然语言处理、金融量化分析，还是计算机视觉，如果你在使用 Python 进行高性能计算，处理矩阵和张量，或包含其他计算密集型运算， Numba 提供的加速效果可以比肩原生的 C/C++ 程序，只需要在函数上添加一行 @jit 的装饰。它支持 CPU 和 GPU ，是数据科学家必不可少的编程利器。

参考文献

① Python程序的执行过程（解释型语言和编译型语言）

② 加一行注释，让你的 Python 程序提速 10+ 倍！Numba 十分钟上手指南