【深度好文】Python加速库Numba简介

1. Numba是什么？

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

2. 如何安装Numba？

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

$ conda install numba

或者

$ pip install numba

3. 初识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直接求和，调用函数np.sum()，代码如下：

import numpy as np
def sum_numpy(M):
    return np.sum(M)

3) numba实现

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

from numba import jit
def 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 = time.time()
    print('current Function [%s] run time is %.2f' % (func.__name__ ,end - start))

if __name__ == "__main__":
    np.random.seed(1234)
    M = np.random.random([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)

运行结果如下：

4. Numba进阶

1）使用@jit 装饰器

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

@jit
def 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__":
    np.random.seed(1234)
    M = np.random.random([int(4e3)] * 2)
    loops = 5
    print_run_time(sum_matrix_jit,M,loops)
    print_run_time(sum_matrix_njit, M, loops)

运行结果如下：

5. 总结

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

6. 参考

参考 链接一 链接二

关注公众号《AI算法之道》，获取更多AI算法资讯。