前言
如果你在使用 Python
进行高性能计算, Numba
提供的加速效果可以比肩原生的 C/C++
程序,只需要在函数上添加一行 @jit
的装饰。它支持 CPU
和 GPU
,是数据科学家必不可少的编程利器。
我们知道,计算机只能执行二进制的机器码,C、C++
等编译型语言依靠编译器将源代码转化为可执行文件后才能运行,Python、Java
等解释型语言使用解释器将源代码翻译后在虚拟机上执行。对于 Python
,由于解释器的存在,其执行效率比 C
语言慢几倍甚至几十倍。
上图比较了当前流行的各大编程语言在几个不同任务上的计算速度。C
语言经过几十年的发展,优化已经达到了极致。以 C
语言为基准,大多数解释语言,如 Python、R
会慢十倍甚至一百倍。 Julia
这个解释语言是个“奇葩”,因为它采用了 JIT
编译技术。
解决 Python
执行效率低的问题,一种解决办法是使用 C/C++
语言重写 Python
函数,但是这要求程序员对 C/C++
语言熟悉,且调试速度慢,不适合绝大多数 Python
程序员。另外一种非常方便快捷的解决办法就是使用 Just-In-Time(JIT)
技术,本文将解释 JIT
技术的原理,并提供几个案例,让你十分钟内学会 JIT
技术。
一、Python解释器的工作原理
Python
是一门解释语言, Python
为我们提供了基于硬件和操作系统的一个虚拟机,并使用解释器将源代码转化为虚拟机可执行的字节码。字节码在虚拟机上执行,得到结果。
我们使用 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
装饰器前后的运行结果为:
测试①:
这是未添加 jit
装饰器时的结果。
测试②:
这是添加了 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)
测试③的结果为:
测试④的结果为:
可以看见两次测试的结果相差无几,为什么呢?
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))
运行结果为:
代码中两次调用 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)
测试⑤和⑥的结果分别为:
@jit(int32(int32, int32))
告知 Numba
你的函数在使用什么样的输入和输出,括号内是输入,括号左侧是输出。这样不会加快执行速度,但是会加快编译速度,可以更快将函数编译到机器码上。
五、Numba到底有多快
网上有很多对 Numba
进行性能评测的文章,在一些计算任务上, Numba
结合 NumPy
,可得到接近 C
语言的速度。
六、Numba的更多功能
除了上面介绍的加速功能,Numba
还有很多其他功能。 @vectorize
装饰器可以将一个函数向量化,变成类似 NumPy
函数一样,直接处理矩阵和张量。R
语言用户可能非常喜欢这个功能。
Numba
还可以使用 GPU
进行加速,目前支持英伟达的 CUDA
和 AMD
的 ROC
。 GPU
的工作原理和编程方法与 CPU
略有不同。
七、Numba原理
Numba
使用了 LLVM
和 NVVM
技术,这个技术可以将 Python、Julia
这样的解释语言直接翻译成 CPU
或 GPU
可执行的机器码。
总结
无论你是在做 NLP
自然语言处理、金融量化分析,还是计算机视觉,如果你在使用 Python
进行高性能计算,处理矩阵和张量,或包含其他计算密集型运算, Numba
提供的加速效果可以比肩原生的 C/C++
程序,只需要在函数上添加一行 @jit
的装饰。它支持 CPU
和 GPU
,是数据科学家必不可少的编程利器。