动手

动手

用Numba实现极速数据处理

作者/Fernando Doglio

Globant公司软件架构师。过去十年一直从事Web开发工作,期间使用了大多数最前沿的技术,如PHP、Ruby on Rails、MySQL、Python、Node.js、AngularJS、REST API等。

Fernando喜欢钻研新事物,他的GitHub账户每个月也会因此获得回购。他还是开源拥护者,并通过网站lookingforpullrequests.com来获得人们的支持。Fernando另著有Pro REST API Development with Node.js。他的Twitter账号是@deleteman123。

数据处理(number crunching)是编程世界的一个主题。但是,由于Python经常用于解决科学研究和数据科学问题,所以数据处理成了Python世界的主流课题。

接下来,我会介绍一种方法Numba,它可以帮助我们写出快速高效的代码以解决科学计算问题。下面,我们从安装开始讲起,然后通过一些示例代码体现方法的优势。

Numba

Numba(http://numba.pydata.org/)是一个模块,让你能够(通过装饰器)控制Python解释器把函数转变成机器码。因此,Numba实现了与C和Cython同样的性能,但是不需要用新的解释器或者写C代码。

这个模块可以按需生成优化的机器码,甚至可以编译成CPU或GPU可执行代码。

下面的代码是官方网站上的示例,可以显示模块的使用方法。我们将在后面详细介绍:

from numba import jit
from numpy import arange

# jit装饰器告诉Numba编译函数
# 当函数被调用时,Numba会把参数类型引入
@jit
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i, j]
    return result

a = arange(9).reshape(3, 3)
print(sum2d(a))

虽然Numba看起来好像非常给力,但是它只是针对数组操作进行优化。它非常适合配合NumPy使用。因此,并非每个函数都可以用Numba优化,滥用Numba甚至会损害性能。

例如,让我们看一个类似的例子,不用Numba也可以完成类似的任务:

from numba import jit
from numpy import arange

# jit装饰器告诉Numba编译函数
# 当函数被调用时,Numba会把参数类型引入
@jit
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i, j]
    return result

a = arange(9).reshape(3, 3)
print(sum2d(a))

前面的代码使用和不使用@jit行的效果如下。

  • 使用@jit行:0.3秒
  • 不使用@jit行:0.1秒

安装

安装Numba有两种方法:可以用Anaconda出品的conda包管理器,也可以复制GitHub项目源代码进行编译。

如果你打算用conda方法安装,需要先安装miniconda命令行工具(可以从http://conda.pydata.org/miniconda.html下载)。安装之后,输入下面的命令:

$ conda install numba

命令输出结果如下图所示。所有要安装、升级的包都会显示出来,像numpyllvmlite都是与Numba有直接依赖的包。

{%}

另外,如果想用源代码安装,你需要先用下面的命令复制源代码:

$ git clone git://github.com/numba/numba.git

当然numpyllvmlite包也是需要提前安装好的。都准备好之后,用下面的命令进行安装:

$ python setup.py build_ext –inplace

需要注意的是,即使没有安装依赖包,上面的命令也可以成功运行。但是如果你没有安装依赖,Numba是没法儿用的。

要检查Numba是否可以正常使用,可以在Python的REPL里输入下面的命令:

>>> import numba
>>> numba.__version__
'0.18.2'

使用Numba

现在Numba已经安装好了,让我们看看如何使用它。这个模块提供的主要功能如下:

  • 即时代码生成(On-the-fly code generation)
  • CPU和GPU原生代码生成
  • 与具有NumPy依赖的Python科学计算软件配合使用

  • Numba代码生成

    Numba代码生成的主要方式是使用@jit装饰器。加上它就表示要用Numba的JIT编译器对函数进行优化。

    让我们看看如何用@jit装饰器进行优化。

    使用这个装饰器的方式有几种。默认的,也是官方推荐的方法,之前也已经介绍过:

    延迟编译(Lazy compilation)
    
    

    在下面的代码中,当函数被调用时,Numba将生成优化代码。它将引用属性类型和函数的返回类型:

    from numba import jit
    
    @jit
    def sum2(a,b):
        return a + b
    
    

    如果你用同样的函数调用其他类型,会生成并优化不同的代码路径。

    (1) 及时编译

    另一方面,如果你知道函数的接收类型(返回类型也可以),可以把这些类型传到@jit装饰器。之后,只有这种特殊情况会被优化。

    下面代码中增加的部分会被传递到函数的签名里:

    from numba import jit, int32
    
    @jit(int32(int32, int32))
    def sum2(a,b):
        return a + b
    
    

    用于指定函数签名的常用类型如下。

    • void:函数返回值类型,表示不返回任何结果。
    • intpuintp:指针大小的整数,分别表示签名和无签名类型。
    • intcuintc:相当于C语言的整型和无符号整型。
    • int8int16int32int64:固定宽度整型(无符号整型前面加u,比如uint8)。
    • float32float64:单精度和双精度浮点数类型。
    • complex64complex128:单精度和双精度复数类型。
    • 数组可以用任何带索引的数值类型表示,比如float32[:]就是一维浮点数数组类型,int32[:,:]就是二维整型数组。

    (2) 其他配置

    除了及时编译,还有两个编译选项可以添加到@jit装饰器上。这两个选项将帮助我们完成Numba优化。选项具体描述如下。

    (a) 没有GIL

    无论何时,只要我们的代码用原始类型优化(不是用Python类型),GIL就不再必要了。

    有一种方法可以禁止GIL。我们可以把nogil=True属性传到装饰器。这样我们就可以用多线程运行Python代码(或者说是Numba代码)了。

    也就是说,只要不再受GIL的限制,你就可以处理多线程系统的常见问题了(一致性、数据同步、竞态条件等)。

    (b) 无Python模式

    这个选项可以让我们设置Numba的编译模式。默认设置时,它将在模式之间跳转。Numba将针对需要优化的代码自动设置对应的优化模式。

    一共有两种模式。一种是object模式。它产生的代码可以处理所有Python对象,并用C API完成Python对象上的操作。另一种是nopython模式,它可以不调用C API而生成更高效的代码。唯一的问题是,只有一部分函数和方法可以使用。

    如果Numba不利用循环JIT(loop-jitting)方法,object模式就不会产生更快的代码(就是说循环可以被提取然后编译成nopython模式)。

    我们可以用Numba把代码强制转换成nopython模式,如果转换失败就会产生错误。可以用下面的代码实现:

        @jit(nopython=True)
        def add2(a, b):
            return a + b
    
    

    nopython模式的问题在于它有一些限制,除了这种模式支持的Python子集范围有限之外,还有:

    • 函数里表示数值的所有原生类型都可以被引用
    • 函数里不可以分配新内存

    另外,由于使用循环JIT方式,被优化的循环内部不能产生返回状态。否则,这种情况不适合优化。

    下面,让我们用一段示例代码来演示优化的过程:

        def sum(x, y):
            array = np.arange(x * y).reshape(x, y)
            sum = 0
            for i in range(x):
                for j in range(y):
                    sum += array[i, j]
            return sum
    
    

    上面的代码取自Numba网站。这个函数适合循环JIT,也叫循环切换(loop-lifting)。为了让代码如预期运行,我们用Python REPL模式:

    {%}

    另外,我们还直接使用了函数的inspect_types方法。这样做的好处是可以看到函数的源代码。这样在匹配Numba生成的机器码时,可以和源代码进行对照。

    上面这个方法的输出结果可以帮助我们理解Numba优化背后的含义。更具体地说,可以看到具体的引用类型,是否进行了自动优化,以及每行Python代码被翻译成多少行代码。

    要理解inspect_types方法从代码中获得的输出结果的含义,需要注意看每一行源代码的编译过程都是由单独的行号开始的。后面跟着的是这一行生成的汇编指令,最后你可以看到不加注释的Python源代码。

    注意看LiftedLoop行。在这一行你会看到Numba的优化代码。还需要注意在许多行最后引用的Numba类型。无论何时你看到一个pyobject类型,都不是原生类型,而是普通Python类型的封装版。

  • 在GPU上运行代码

    前面已经提过,Numba代码可以运行在CPU和GPU上。其实,在GPU上运行并行程序可以进一步提升性能,比CPU上运行更快。

    更具体地说就是,Numba支持CUDA编程(http://www.nvidia.com/object/cuda_home_new.html),按照CUDA模式的规则把一部分Python代码翻译成CUDA核心与设备支持的语言形式。

    CUDA是Nvidia开发的并行计算平台和编程模式。它可以利用GPU获得更大的速度提升。

    因为GPU编程是个较大的主题,可以写一整本书,所以这里不再介绍更多细节。我们只说明Numba具有这种能力,通过装饰器@cuda.jit就可以实现。关于这个主题的具体文档,请参考http://numba.pydata.org/numba-doc/0.18.2/cuda/index.html的内容。

 

{%}

《Python性能分析与优化》对于Python程序员来说,仅仅知道如何写代码是不够的,还要能够充分利用关键代码的处理能力。本书将讨论如何对Python代码进行性能分析,找出性能瓶颈,并通过不同的性能优化技术消除瓶颈。