The website uses cookies. By using this site, you agree to our use of cookies as described in the Privacy Policy.
I Agree

Numba(1) —— 加速你的Python

Python是一门优雅、简洁、跨平台的编程语言。最为人诟病的一点是Python实在是太慢了、太慢了,几乎找不到一门动态语言能有Python那般蜗牛的计算速度。

对于I/O密集型的任务,Python可以启用多线程来提高速度。而对于CPU密集型(计算密集型)的任务。对于科学计算的任务,Python跟MATLAB相比速度差的不是一点点。

一种提升Python计算速度的方式是改用numpy,把标准库math中的函数用numpy中同类型的函数替代能起到一定的加速作用;把列表转换成np.ndarray使用numpy的操作函数比使用列表的方法要快。

另一种更能加速Python计算速度的方法就是合理的使用numba

numba

Numba是Python的即时编译器,在使用NumPy数组和函数以及循环的代码上效果最佳。使用Numba的最常见方法是通过其装饰器集合,这些装饰器可应用于您的函数以指示Numba对其进行编译。调用带有Numba装饰的函数时,该函数将被“即时”编译为机器代码以执行,并且您的全部或部分代码随后可以本机机器速度运行!—— numba 官方文档。

使用numba最简单的方式就是在计算密集的函数前面加上@jit修饰器

numba在如下所示的代码上运行良好:

from numba import jit
import numpy as np

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

@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
def go_fast(a): # Function is compiled to machine code when called the first time
    trace = 0
    for i in range(a.shape[0]):   # Numba likes loops
        trace += np.tanh(a[i, i]) # Numba likes NumPy functions
    return a + trace              # Numba likes NumPy broadcasting

print(go_fast(x))

请注意,numba无法理解pandas,因此numba只需通过解释器运行此代码,但会增加numba内部开销!

类似的,在 @jit修饰的函数内部,不建议使用内置的函数诸如sum,而建议改用numpy.sum,因为numba无法识别内置的sum函数,而对numpy.sum支持良好。

nopython 模式

numba @jit装饰器基本上在两种编译模式下运行,即nopython模式和对象模式。在上面的go_fast示例中,在@jit装饰器中设置了nopython = True,这指示numbanopython模式下运行。nopython编译模式的行为是实质上编译装饰后的函数,以便其完全运行而无需Python解释器的参与。这是使用numba jit装饰器的推荐和最佳实践方法,因为它可以带来最佳性能。

如果nopython模式下的编译失败,numba可以使用对象模式进行编译,如果未设置nopython = True,则这是@jit装饰器的后备模式。在这种模式下,numba将识别可以编译的循环并将其编译为在机器代码中运行的函数,并在解释器中运行其余代码。为了获得最佳性能,请避免使用此模式!

numba的性能

首先,回想一下,numba必须在执行函数的机器代码版本之前针对给定的参数类型编译函数。但是,一旦编译完成,numba会为呈现的特定类型的参数缓存函数的机器代码版本。如果使用相同类型再次调用,则它可以重用缓存的版本,而不必再次编译。

衡量性能时,一个真正常见的错误是不考虑上述行为,并使用一个简单的计时器对代码进行一次计时,该计时器包括在执行时间中编译函数所花费的时间。

以下是一个测量numba.jit修饰的函数性能的例子

from numba import jit
import numpy as np
import time

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

@jit(nopython=True)
def go_fast(a): # Function is compiled and runs in machine code
    trace = 0
    for i in range(a.shape[0]):
        trace += np.tanh(a[i, i])
    return a + trace

# DO NOT REPORT THIS... COMPILATION TIME IS INCLUDED IN THE EXECUTION TIME!
start = time.time()
go_fast(x)
end = time.time()
print("Elapsed (with compilation) = %s" % (end - start))

# NOW THE FUNCTION IS COMPILED, RE-TIME IT EXECUTING FROM CACHE
start = time.time()
go_fast(x)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start))

输出结果为:

Elapsed (with compilation) = 0.33030009269714355
Elapsed (after compilation) = 6.67572021484375e-06

numba的其他修饰器

numba有很多装饰器,我们已经看到了@jit,但是还有:

  • @njit
  • @vectorize
  • @guvectorize
  • @stencil
  • @jitclass
  • @cfunc
  • @overload

此外,numba还支持调用GPU

Measure
Measure
Summary | 2 Annotations
numba无法理解pandas,因此numba只需通过解释器运行此代码,但会增加numba内部开销!
2021/02/07 12:04
执行函数的机器代码版本之前针对给定的参数类型编译函数
2021/02/07 12:05