Python 提速大杀器之 numba 篇

OpenMMLab

小助手：OpenMMLabwx 微信公众号同名

文@132596

你是不是曾经有这样的苦恼，python 真的太好用了，但是它真的好慢啊（哭死) ; C++ 很快，但是真的好难写啊，此生能不碰它就不碰它。老天啊，有没有什么两全其美的办法呢？俗话说的好：办法总是比困难多，大家都有这个问题，自然也就有大佬来试着解决这个问题，这就请出我们今天的主角: numba

不过在介绍 numba 之前，我们还是得来看看 python 为什么这么慢：

为什么 python 这么慢

用过 python 的人都知道，尤其是在有循环的情况下，python 会比 C++ 慢很多，所以很多人都避免在 python 代码里引入复杂的 for 循环。我们可以想想 python 和 C++ 写起来有哪些区别呢：

动态变量

如果你写过 C/C++ 就会发现，我们需要对变量类型有严格的定义，我们需要定义变量的类型是 int 或者 float 之类的。但是 python 就不一样了，写过的 python 的人都知道，它去掉了变量申明和数据类型。也就是说，无论啥数据，咱啥都不用管，想存就存！那么 python 是如何做到这样洒脱自由的呢？这就不得不提 python 中万物皆是对象了，真正的数据是存在对象里面的。对于一个简单的两个变量的加法，python 每次在做运算的时候都得先判断变量的类型，再取出来进行运算，而对于 C 来说，简单的内存读写和机器指令 ADD 即可。其实在 C/C++ 中也有可变数据类型，但是其声明是非常复杂的，是一种非常令人头疼的结构。

解释性语言

C/C++ 这类编译性语言最大的好处就是其编译过程是发生在运行之前的，源代码在调用前被编译器转换为可执行机器码，这样就节约了大量的时间。而 python 作为一种解释性语言，没法做到一次编译，后续可以直接运行，每次运行的时候都要重新将源代码通过解释器转化为机器码。这样一个好处就是非常容易 debug（这里要再次感叹一下 python 真不愧是新手友好型语言~），当然，这个问题自然也是有尝试解决的办法，一个很重要的技术就是 JIT （ Just-in-time compilation ）：JIT 即时编译技术是在运行时（runtime）将调用的函数或程序段编译成机器码载入内存，以加快程序的执行。说白了，就是在第一遍执行一段代码前，先执行编译动作，然后执行编译后的代码。

上面只是简单列出了两点，当然还有更多的原因，限于篇幅就不再具体介绍，而我们开篇提到的 numba 就是通过 JIT 加速了 python 代码。那么怎么使用 numba 加速我们的代码呢？我们可以看一些简单的例子：

numba 加速 python 的小例子

用 numba 加速 python 代码多简单方便呢，我们先来看看如何使用 numba 加速 python 代码：

如果让你用单纯的 python 计算一个矩阵所有元素的和，很容易可以写出下面的代码：

def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result

当需要计算的矩阵很小的时候，貌似速度也不慢，可以接受，但是如果输入的矩阵大小为 (500, 500)，

a = np.random.random((500, 500)) 
%timeit cal_sum(a)

输出结果为：

47.8 ms ± 499 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我们尝试加上 numba：

import numba  
@numba.jit(nopython=True) 
def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result

输入同样大小的矩阵

a = np.random.random((500, 500)) 
%timeit cal_sum(a)

输出结果为：

236 µs ± 545 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

注意在这里我们使用了 %itemit 测试运行时间（原因我们留到后面说），通过对比两个时间，我们可以发现通过 numba 获得了非常明显的加速效果！

我们来具体看一下如何用 numba 加速 python 代码：在实际使用过程中，numba 其实是以装饰器的形式加在 python 函数上的，用户可以不用关心到底 numba 是通过什么方法来优化代码，只需要调用就行。同时需要注意到 @jit 装饰器同时也有一个参数 nopython，这个参数主要是来区分 numba 的运行模式，numba 其实有两种运行模式：一个是 nopython 模式，另一个就是 object 模式。只有在 nopython 模式下，才会获得最好的加速效果，如果 numba 发现你的代码里有它不能理解的东西，就会自动进入 object 模式，保证程序至少是能够运行的（当然这其实就失去了添加 numba 的意义）。如果我们将装饰器改为 @jit(nopython=True) 或者 @njit ，numba 会假设你已经对所加速的函数非常了解，强制使用加速的方式，不会进入 object 模式，如编译不成功，则直接抛出异常。

当然说到这里，可能大家还是很困惑，numba 到底是怎么加速 python 代码的？

python 代码的编译过程包括四个阶段： 词法分析 -> 语法分析 -> 生成字节码 -> 将字节码解释为机器码执行, 常见的 python 解释器的类型有 cpython、IPython、PyPy、Jython、IronPython，与其他解释器不同， numba 是使用 LLVM 编译技术来解释字节码的。

LLVM 是一个编译器，它采用字节码，并将其编译为机器码，编译过程涉及许多额外的传递，而 LLVM编译器可以优化字节码，例如某些频繁执行的模块，LLVM 可以将其作为 “hot code” 从而进行相应的优化，LLVM 工具链非常擅长优化字节码，它不仅可以编译 numba 的代码，还可以优化它。

在第一次调用 numba 装饰的函数时，numba 将在调用期间推断参数类型，numba 会结合给定的参数类型将其编译为机器代码。这个过程是有一定的时间消耗的，但是一旦编译完成，numba 会为所呈现的特定类型的参数缓存函数的机器代码版本，如果再次使用相同的类型调用它，它可以重用缓存的机器代码而不必再次编译。

在测量性能时，如果只使用一个简单的计时器来计算一次，该计时器包括在执行时编译函数所花费的时间，最准确的运行时间应该是第二次及以后调用函数的运行时间。
对于指定输入类型这个问题，我们可以尝试做一个简单的实验看看到底有怎样的影响：

a = np.random.random((5000, 5000)) 
# 第一次调用时间包括编译时间 
start = time.time() 
cal_sum(a) 
end = time.time() 
print("Elapsed (with compilation) = %s" % (end - start)) 
# 函数被编译，机器代码被缓存 
start = time.time() 
cal_sum(a) 
end = time.time() 
print("Elapsed (after compilation) = %s" % (end - start)) 
# 这里 a 本身的类型为 np.float64 
b = a.astype(np.float32) 
# 调用相同的函数，但是输入数据的类型变为 np.float32 
start = time.time() 
cal_sum(b) 
end = time.time() 
print("Elapsed (after compilation) = %s" % (end - start))

输出结果：

Elapsed (with compilation) = 0.20406198501586914 
Elapsed (after compilation) = 0.025263309478759766 
Elapsed (after compilation) = 0.07892274856567383

可以看到如果我们输入了和第一次调用编译时不同的数据类型，函数的运行时间也会有一个很明显的增加，但仍然是远低于第一次运行时的编译的时间。

3. 如果调用 numba 的时候 显式地指定输入、输出数据的类型，可以加快初次调用的函数时的编译速度， 同时坏处就是如果显式指定后，那么之后调用该函数都必须满足规定的数据类型。

a = np.random.random((500, 500)) 
@numba.njit() 
def cal_sum1(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result 
@numba.njit('float64(float64[:, :])') 
def cal_sum2(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result 
# 不指定输入输出数据类型，让 numba 自己判断 
start = time.time() 
cal_sum1(a) 
end = time.time() 
print("Elapsed (with compilation) = %s" % (end - start)) 
# 指定输入输出数据类型 
start = time.time() 
cal_sum2(a) 
end = time.time() 
print("Elapsed (with compilation) = %s" % (end - start))

分别耗时：

Elapsed (after compilation) = 0.054465532302856445 
Elapsed (after compilation) = 0.0004112720489501953

可以看到编译的时间被大大减少了，其实这个时间非常接近直接运行该函数生成的机器代码的时间。

上面说了这么多，但是转念一想，矩阵相加这个函数 numpy 里好像早就有了， np.sum 它不好用，它不香嘛？？干嘛搞得这么复杂？

好吧，就上面举的简单的例子来说，使用 numpy 和 numba 加速基本效果差不多，但是在实际情况里面，不是所有的 for 循环代码都可以直接用 numpy 自带的函数实现。 但是 numba 基本对所有的 for 循环代码都有非常好的加速效果，当然前提是 for 循环里面的代码必须是 numba 能够理解的。

而在从实际使用中，一般推荐将代码中密集的计算部分提取出来作为单独的函数实现，并使用 nopython 方式优化，这样可以保证我们能使用到 numba 的加速功能。其余部分还是使用 python 原生代码，这样一方面就可以做到在 numba 加速不明显或者无法加速的代码中调用各种函数实现自己的代码逻辑，另一方面也能享受到 numba 的加速效果。

numba 加速 numpy 运算

上面说了 numba 一大亮点就是加速 for 循环，除此以外，numba 对 numpy 的运算也同样的有加速的效果。因为即使是 numpy 也没有 numba 转换为机器码快，numba 尤其擅长加速 numpy 的基本运算 (如加法、相乘和平方等等) ，其实准确来说如果 numpy 函数是对各个元素采用相同的操作的情况下，都会有比较好的效果。我们简单举一个 numba 加速 numpy 运算的例子：

a = np.ones((1000, 1000), np.int64) * 5 
b = np.ones((1000, 1000), np.int64) * 10 
c = np.ones((1000, 1000), np.int64) * 15 
def add_arrays(a, b, c): 
    return np.square(a, b, c) 
@numba.njit 
def add_arrays_numba(a, b, c): 
    return np.square(a, b, c) 
# 第一次调用完成编译 
add_arrays_numba(a) 
# 函数被编译，机器代码被缓存 
start = time.time() 
add_arrays_numba(a) 
end = time.time() 
print("Elapsed (after compilation) = %s" % (end - start)) 
# 不使用 numba 加速 
start = time.time() 
add_arrays(a) 
end = time.time() 
print("Elapsed = %s" % (end - start))

Elapsed (after compilation) = 0.002088785171508789
Elapsed = 0.0031290054321289062

当我们对 numpy 数组进行基本的数组计算，比如加法、乘法和平方，numpy 都会自动在内部向量化，这也是它可以比原生 python 代码有更好性能的原因。但是在特定情况下，numpy 的代码也不会和优化过的机器代码速度一样快，此时 numba 直接作用于 numpy 运算也能起到一定的加速效果。

另一个例子主要来自于 MMDetection3D，经过一定的简化，主要是用来计算将点的坐标 (x, y) 压缩到给定的 [x_min, y_min, x_max, y_max] 范围内:

x = np.random.random((5000))*5000 
y = np.random.random((5000))*5000 
x_min = 0 
x_max = 1000 
y_min=0 
y_max=2000 
@numba.njit 
def get_clip_numba(x, y, x_min, y_min, x_max, y_max): 
    z = np.stack((x, y), axis=1) 
    z[:, 0] = np.clip(z[:, 0], x_min, x_max) 
    z[:, 1] = np.clip(z[:, 1], y_min, y_max) 
    return z 
def get_clip(x, y, x_min, y_min, x_max, y_max): 
    z = np.stack((x, y), axis=1) 
    z[:, 0] = np.clip(z[:, 0], x_min, x_max) 
    z[:, 1] = np.clip(z[:, 1], y_min, y_max) 
    return z 
%timeit get_clip_numba(x, y, x_min, y_min, x_max, y_max) 
%timeit get_clip(x, y, x_min, y_min, x_max, y_max)

分别用时：

33.8 µs ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 
57.2 µs ± 258 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

从实际情况来看，并不是所有的 numpy 函数在使用 numba 后都能获得比较好的加速效果，在某些情况下甚至会降低 numpy 的运行速度。因此，在实际使用过程中建议提前测试一下确认加速效果。通常将 numba 用于加速 numpy 的时候都是 for 循环和 numpy 一起使用的情况。 numba 对 numpy 的大部分常用的函数都做了支持。

numba 使用 CUDA 加速

numba 更厉害的地方就在于，我们可以直接用 python 写 CUDA Kernel，直接在 GPU 上编译和运行我们的 Python 程序，numba 通过将 python 代码直接编译为遵循 CUDA 执行模型的 CUDA 内核和设备函数来支持 CUDA GPU 编程（但是实际上 numba 目前支持的 CUDA API 很少，希望开发团队能更肝一点~~~) ，对于不熟悉 CUDA 的同学，我们推荐大家看一下之前的一篇文章 PyTorch 源码解读之 cpp_extension：揭秘 C++/CUDA 算子实现和调用全流程了解一下 CUDA 编程的基本概念。为了节省将 numpy 数组复制到指定设备，然后又将结果存储到 numpy 数组中所浪费的时间，numba 提供了一些函数来声明并将数组送到指定设备来节省不必要的复制到 cpu 的时间。

常用内存分配函数：

cuda.device_array() ：在设备上分配一个空向量，类似于 numpy.empty() ；
cuda.to_device() ：将主机的数据拷贝到设备；
cuda.copy_to_host() ：将设备的数据拷贝回主机；

我们可以通过一个简单的矩阵相加的例子来看看通过 numba 使用 CUDA 加速的效果：

from numba import cuda # 从numba调用cuda
import numpy as np
import math
from time import time
@cuda.jit
def matrix_add(a, b, result, m, n):
    idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    idy = cuda.threadIdx.y+ cuda.blockDim.y * cuda.blockIdx.y
    if idx < m and idy < n:
        result[idx, idy] = a[idx, idy] + b[idx, idy]
m = 5000
n = 4000
x = np.arange(m*n).reshape((m,n)).astype(np.int32)
y = np.arange(m*n).reshape((m,n)).astype(np.int32)
# 拷贝数据到设备端
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)
# 在显卡设备上初始化一块用于存放GPU计算结果的空间
gpu_result1 = cuda.device_array((m,n))
gpu_result2 = cuda.device_array((m,n))
cpu_result = np.empty((m,n))
threads_per_block = 1024
blocks_per_grid = math.ceil(m*n / threads_per_block)
# 第一次调用包含编译时间
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result1, m, n)
cuda.synchronize()
print("gpu matrix add time (with compilation) " + str(time() - start))
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result2, m, n)
cuda.synchronize()
print("gpu matrix add time (after compilation)" + str(time() - start))
start = time()
cpu_result = np.add(x, y)
print("cpu matrix add time " + str(time() - start))

运行时间分别为：

gpu matrix add time (with compilation) 0.15977692604064941
gpu matrix add time (after compilation) 0.0005376338958740234
cpu matrix add time 0.023023128509521484

在通过 numba 进行 CUDA 加速的时候，主要是通过调用 @cuda.jit 装饰器实现，从结果可以看到 numba 通过调用 CUDA 明显加速了 python 程序。

For 循环写法的影响

下面的一段代码截取自 MMDetection3D, 主要是用来判断一系列点是否在一系列多边形的内部，我们可以有如下的两种写法：

在 For 循环里面计算 vec1，每次循环都需要访问多边形 polygon 变量

@numba.jit(nopython=True) 
def points_in_convex_polygon1(points, polygon, clockwise=True): 
    # first convert polygon to directed lines 
    num_points_of_polygon = polygon.shape[1] 
    num_points = points.shape[0] 
    num_polygons = polygon.shape[0] 
    vec1 = np.zeros((2), dtype=polygon.dtype) 
    ret = np.zeros((num_points, num_polygons), dtype=np.bool_) 
    success = True 
    cross = 0.0 
    for i in range(num_points): 
        for j in range(num_polygons): 
            success = True 
            for k in range(num_points_of_polygon): 
                if clockwise: 
                    vec1 = polygon[j, k] - polygon[j, k - 1] 
                else: 
                    vec1 = polygon[j, k - 1] - polygon[j, k] 
                cross = vec1[1] * (polygon[j, k, 0] - points[i, 0]) 
                cross -= vec1[0] * (polygon[j, k, 1] - points[i, 1]) 
                if cross >= 0: 
                    success = False 
                    break 
            ret[i, j] = success 
    return ret

在循环前预先计算好所有的 vec

@numba.jit(nopython=True) 
def points_in_convex_polygon2(points, polygon, clockwise=True): 
    # first convert polygon to directed lines 
    num_points_of_polygon = polygon.shape[1] 
    num_points = points.shape[0] 
    num_polygons = polygon.shape[0] 
    # vec for all the polygons 
    if clockwise: 
        vec1 = polygon - polygon[:, np.array([num_points_of_polygon - 1] + 
                                 list(range(num_points_of_polygon - 1))), :] 
    else: 
        vec1 = polygon[:, np.array([num_points_of_polygon - 1] + 
                       list(range(num_points_of_polygon - 1))), :] - polygon 
    ret = np.zeros((num_points, num_polygons), dtype=np.bool_) 
    success = True 
    cross = 0.0 
    for i in range(num_points): 
        for j in range(num_polygons): 
            success = True 
            for k in range(num_points_of_polygon): 
                vec = vec1[j,k] 
                cross = vec[1] * (polygon[j, k, 0] - points[i, 0]) 
                cross -= vec[0] * (polygon[j, k, 1] - points[i, 1]) 
                if cross >= 0: 
                    success = False 
                    break 
            ret[i, j] = success 
    return ret

简单测试一下两种写法的速度：

points = np.random.random((20000, 2)) * 100 
polygon = np.random.random((1000, 100, 2)) * 200  
start = time.time() 
points_in_convex_polygon1(points, polygon) 
end = time.time() 
print("Elapsed (with compilation) = %s" % (end - start)) 
start = time.time() 
points_in_convex_polygon1(points, polygon) 
end = time.time() 
print("Elapsed (after compilation) = %s" % (end - start)) 
start = time.time() 
points_in_convex_polygon2(points, polygon) 
end = time.time() 
print("Elapsed (with compilation) = %s" % (end - start)) 
start = time.time()