Python加速库Numba的安装和常见使用方法

Numba是一款可以将python函数编译为机器代码的JIT编译器，经过Numba编译的python代码（仅限数组运算），其运行速度可以接近C或FORTRAN语言。普通python语言靠CPython编译的，但是Numba使用Jit编译器可以直接将一个函数转化为机器码。

01

—

安装

我们可以通过conda或者pip来进行安装, 命令如下：

$ conda install numba

或者

$ pip install numba

02

—

初识Numba

为了加强大家对Numba库的直观了解，我们首先以求二维数组元素之和作为例子来讲解。 1) 常规实现 为了求二维数据元素之和，最简单直接的实现如下所示：

def sum_matrix(M):    N1, N2 = M.shape    result = 0    for i in range(N1):        for j in range(N2):            result += M[i, j]    return result

2) numpy实现

作为对比，可以使用numpy直接求和，调用函数()，代码如下：

import numpy as npdef sum_numpy(M):    return (M)

3) numba实现 尽管numpy库有实现二维数组求和的现成函数 ，使用numba中的jit可以将函数 sum_matrix运行的和一样快。只需要将sum_matrix作为参数传递给jit即可。

from numba import jitdef sum_jit(M):    sum_matrix_jit = jit()(sum_matrix)    return sum_matrix_jit

 4）性能对比 为了统计上述函数运行时间，我们统计三个函数的运行时间如下所示：

def print_run_time(func,M,loops):    start = time.time()    for _ in range(loops):        func(M)    end = ()    print('current Function [%s] run time is %.2f' % (func.__name__ ,end - start))
if __name__ == "__main__":    np.random.seed(1234)    M = ([int(4e3)] * 2)    loops = 5    print_run_time(sum_matrix,M,loops)    print_run_time(sum_numpy,M,loops)    print_run_time(sum_jit,M,loops)

运行结果如下：

03

—

Numba进阶

 1）使用@jit 装饰器

使用jit来转换函数，一种更为常见的方式是使用装饰器的写法，如下所示：

@jitdef sum_matrix_jit(M):    N1, N2 = M.shape    result = 0    for i in range(N1):        for j in range(N2):            result += M[i, j]    return result

2）使用nopython and @njit

如果希望使用jit加速后的函数完全不受Python解释器的干涉，我们可以使用nopython或者@njit，这种模式下将会提供一个最优的性能来优化for循环。如果代码涉及非数值处理情形，python解释器是重要的，必须避免使用该模式。代码实现如下：

from numba import njit@njit # or @jit(nopython=True)def sum_matrix_njit(M):    N1, N2 = M.shape    result = 0    for i in range(N1):        for j in range(N2):            result += M[i, j]    return result

3）性能对比 为了对比使用@jit和@njit，参考上述例子二者的对比实验如下：

if __name__ == "__main__":    (1234)    M = ([int(4e3)] * 2)    loops = 5    print_run_time(sum_matrix_jit,M,loops)    print_run_time(sum_matrix_njit, M, loops)

运行结果如下：

04

—

总结

本文简单介绍了numba库的安装和常见的使用方法，针对数值运算的python函数，只需要添加简单的@jit装饰器，就可以让运行速度大幅提升，强烈推荐大家在日常的工作学习中使用。