Python 提速大杀器之 numba 篇
时间：2022-10-16 19:30:00
你曾经有过这样的苦恼吗？python 真的很好用，但是真的很慢(哭死) ; C 很快，但写起来真的很难。如果你这辈子不碰它，你就不会碰它。上帝，有什么两全其美的办法吗？俗话说：方法总是比困难多。每个人都有这个问题。当然，也有老板试图解决这个问题。这将邀请我们今天的主角：numba
每次一个Python小技巧！
不过在介绍 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 很少，希望开发团队能更肝一点~~~) ，为了节省将 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()
points_in_convex_polygon2(points, polygon)
end = time.time()
print("Elapsed (after compilation) = %s" % (end - start)) 输出时间：
Elapsed (with compilation) = 3.9232356548309326
Elapsed (after compilation) = 3.6778993606567383
Elapsed (with compilation) = 0.6269152164459229
Elapsed (after compilation) = 0.22288227081298828 通过测试我们可以发现第二种方案会更快，在实际使用的时候，我们可以尽量减少在 for 循环内部内存的访问次数，从而降低函数的运行时间。
总结
我们介绍了一些用 numba 加速的常见场景，能够有效地提高我们代码的速度。不过大家在使用的时候，建议多多尝试，比较一下使用与不使用的速度区别（有时候用了 numba 还可能变得更慢......)，此外 MMDetection3D 很早就使用了 numba 加速代码，而且我们最近在 MMDetection3D 中升级了 numba 的版本，从而获得更好的 numpy 兼容性和代码加速效果，欢迎来玩！
最后祝大家天天进步！！学习Python最重要的就是心态。我们在学习过程中必然会遇到很多难题，可能自己想破脑袋都无法解决。这都是正常的，千万别急着否定自己，怀疑自己。如果大家在刚开始学习中遇到困难，想找一个python学习交流环境，可以 加入我们，领取学习资料，一起讨论。
