资讯详情

Python 提速大杀器之 numba 篇

你曾经有过这样的苦恼吗？python 真的很好用，但是真的很慢(哭死) ; C 很快，但写起来真的很难。如果你这辈子不碰它，你就不会碰它。上帝，有什么两全其美的办法吗？俗话说：方法总是比困难多。每个人都有这个问题。当然，也有老板试图解决这个问题。这将邀请我们今天的主角：

每次一个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 真的不愧是新手友好语言~）， 当然，这个问题自然决这个问题。一项非常重要的技术是说白了，就是在第一次执行代码之前，先执行编译动作，再执行编译后的代码。

以上只列出了两点，当然还有更多的原因，仅限于空间不再具体介绍，我们在开头提到 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)

请注意，我们在这里使用它通过比较两个时间，我们可以发现测试运行时间(原因留在后面)。 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，与其它解释器不同，

LLVM 它是一个使用字节码并将其编译成机器代码的编译器，编译过程涉及许多额外的传输， LLVM编译器可以优化字节码，，编译器可以优化字节码，LLVM 可以作为 “hot code” 从而进行相应的优化，LLVM 工具链非常擅长优化字节码，它不仅可以编译 numba 的代码，它也可以优化。

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

  1. 在测量性能时，如果只使用一个简单的计时器计算一次，则计时器包括在执行过程中编译函数所花费的时间。最准确的操作时间应为第二次和以后调用函数的操作时间。
  2. 对于指定的输入类型，我们可以尝试做一个简单的实验，看看它有什么影响：
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 里好像早就有了， 它不好用，它不香嘛？？干嘛搞得这么复杂？

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

而在从实际使用中，，这样一方面就可以做到在 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学习交流环境，可以 加入我们，领取学习资料，一起讨论。

 

标签： 交流变送器参数

锐单商城拥有海量元器件数据手册、 IC替代型号，打造 电子元器件IC百科大全！

锐单商城 - 一站式电子元器件采购平台