65.9K
CodeProject 正在变化。 阅读更多。
Home

使用 Numba 在 Python 中实现并行

Intel
starIconstarIconstarIconstarIconstarIcon

5.00/5 (2投票s)

2019年10月7日

CPOL
viewsIcon

16667

在本文中，我们将探讨如何通过 Numba 实现并行。

在 Python* 中实现并行对于许多开发者来说一直是一个挑战。在《并行宇宙》第 35 期中，我们探讨了 Python 语言的基础以及实现并行的方式。在本文中，我们将探讨如何通过 Numba* 实现并行。

在 Python 中高效实现并行有三种主要方法：

  1. 通过 Python 的 ctypes 或 cffi 将代码分派到您自己的原生 C 代码（将 C 代码封装在 Python 中）。
  2. 依赖使用高级原生运行时的库，例如 NumPy 或 SciPy。
  3. 使用一个框架，该框架充当引擎，从 Python 或符号数学表达式生成原生速度的代码。

这三种方法都绕过了全局解释器锁 (GIL)，并且是以 Python 社区接受的方式实现的。Numba 框架属于第三种方法，因为它使用即时 (JIT) 和低级虚拟机 (LLVM) 编译引擎来创建原生速度的代码。

使用 Numba 的第一个要求是，您用于 JIT 或 LLVM 编译优化的目标代码必须包含在函数中。在 Python 解释器进行第一次转换（转换为字节码）后，Numba 将查找将函数作为 Numba 解释器通道目标的装饰器。接下来，它将运行 Numba 解释器以生成中间表示 (IR)。之后，它将为目标硬件生成一个上下文，然后继续进行 JIT 或 LLVM 编译。Numba IR 从堆栈机表示转换为寄存器机表示，以实现更好的运行时优化。从此，选项和并行指令的范围将大大扩展。

在下面的示例中，我们使用纯 Python 代码，以便 Numba 有最好的优化机会，而无需指定指令。

import array
import random
from numba import jit
a = array.array('1', [random.randint(0,10) for x in range (0,10000000)])

@jit(nopython=True, parallel=True)
def ssum(x):
   total = 0
   for items in x:
       total+=items
   return total

%timeit sum(a)
111 ms ± 861 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit ssum(a)
4.2 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# Nearly 26X faster!

与混合 CPython 和 NumPy 代码相比，纯 CPython 字节码更容易被 Numba 解释器处理。@jit 装饰器告诉 Numba 在运行函数之前创建 IR，然后创建一个已编译的版本。请注意装饰器上的 nopython 属性。这意味着我们不希望 Numba 无法转换代码时回退到标准的解释器行为（稍后会详细介绍）。我们使用 Python 数组而不是列表，因为它们更容易编译到 Numba。我们还创建了一个自定义的求和函数，因为 Python 的标准 sum 函数具有特殊的迭代器属性，无法在 Numba 中编译。

之前的示例适用于通用的 Python 代码。但是，如果您的代码需要使用 NumPy 或 SciPy 等科学或数值包怎么办？以计算电路的电阻-电容 (RC) 时间常数的以下代码为例。

import numpy as np
test_voltages = np.random.rand(1,1000)*12
test_constants = np.random.rand(1,1000)

def filter_time_constant(voltage, time_constant):
    return voltage * (1-np.exp(1/time_constant))
    
%timeit filter_time_constant(test_voltages, test_constants)
11.2 µs ± 145 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

在这种情况下，由于 NumPy 对 ufunc 的实现，我们将使用 @vectorize 装饰器而不是 @jit。

from numba import vectorize
@vecotrize
def v_filter_time_constant(voltage, time_constant):
    return voltage * (1-np.exp(1/time_constant))
%timeit filter_time_constant(test_voltages, test_constants)
4.74 µs ± 46.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# Over 2x faster!

在处理 NumPy 和 SciPy 等专用框架时，Numba 不仅处理 Python，还处理 NumPy/SciPy 堆栈中的一种特殊类型的原始类型，称为 ufunc，这通常意味着需要使用 C 代码创建 NumPy ufunc，这是一个棘手的任务。在这种情况下，np.exp() 是一个不错的选择，因为它是一个超越函数，并且可以与 Numba 结合使用 Intel® 编译器的高效向量数学库 (SVML) 进行目标定位。@vectorize 和 @guvectorize 都可以使用 Intel 的 SVML 库并帮助处理 NumPy ufunc。

虽然 Numba 具有良好的 ufunc 覆盖率，但重要的是要理解，并非所有 NumPy 或 SciPy 代码库都能在 Numba 中得到很好的优化。这是因为一些 NumPy 原始类型已经得到了高度优化。例如，numpy.dot() 使用了基本的线性代数子程序 (BLAS)，这是一个针对线性代数优化的 C API。如果使用 Numba 解释器，它实际上会产生一个更慢的函数，因为它无法进一步优化 BLAS 函数。为了在 Numba 中最优地使用 ufunc，我们需要寻找一个堆叠的 NumPy 调用，其中许多对数组或向量的操作被复合在一起。例如：

%timeit np.exp(np.arcsin(np.random.rand(1000)))
19.6 µs ± 85.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

@jit(nopython=True)
def test_func(size)
    np.exp(np.arcsin(np.random.rand(1000)))

%timeit test_func(1000)
16 µs ± 80.2 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

@jit 的 Numba 性能略优于直接的 NumPy 代码，因为此计算包含三个而不是一个 NumPy 计算。Numba 可以比 NumPy 本身更好地分析 ufunc 并检测最佳矢量化和对齐。

调整 Numba 的编译指令和性能的另一个领域是使用高级编译选项。主要使用的选项是 nopython、nogil、cache 和 parallel。使用 @jit 装饰器时，Numba 会尝试选择最佳方法来优化给定的代码。但是，如果您对代码的性质有更深入的了解，则可以直接指定编译指令。

第一个选项是 nopython，它防止编译回退到 Python 对象模式。如果代码无法转换，它将向用户抛出错误。第二个选项是 nogil，它在不处理非对象代码时释放 GIL。此选项假定您已经考虑了多线程注意事项，例如一致性和竞态条件。cache 选项将编译后的函数存储在基于文件的缓存中，以避免在下次 Numba 调用同一函数时进行不必要的编译。parallel 指令是对已知可靠的原始类型（如数组和 NumPy 计算）进行的 CPU 定制转换。此选项是符号数学内核的良好首选。

更严格的函数签名提高了 Numba 优化代码的机会。在签名中定义每个参数的预期数据类型，使 Numba 解释器能够获得必要的信息来查找内核的最佳机器表示和内存对齐。这类似于为 C 编译器提供静态类型。以下示例演示了如何向 Numba 提供类型信息。

@jit(int32(int32, int32))
# Expecting int32 values when being processed

@jit([(int64[:], int64 int64[:])]
# Expecting int64 arrays values when being processed

@vectorize([float64([float64(flat64, float64)])
# Expecting float64 arrays values when being processed

总的来说，通过 Numba 在 Python 中访问并行性，关键在于了解一些基本知识，并在积极编写 Python 代码时将这些方法纳入工作流程。以下是该过程的步骤：

  1. 确保核心内核的抽象是合适的。 Numba 要求优化目标位于函数中。不必要的复杂代码可能导致 Numba 编译回退到对象代码。
  2. 查找代码中以某种形式的循环处理数据并具有已知数据类型的地方。 例如，遍历整数列表的 for 循环，或在纯 Python 中处理数组的算术计算。
  3. 如果您使用 NumPy 和 SciPy，请查看可以合并到单个语句中的计算，而不是 BLAS 或 LAPACK 函数。 这些是使用 Numba 的 ufunc 优化功能的绝佳候选。
  • 试验 Numba 的编译选项。
  • 确定函数的预期数据类型签名和核心代码。如果已知（例如 int8 或 int32），则告知 Numba 它应该期望哪些输入数据类型参数。

通过 Numba 实现并行性只需要一些练习和掌握正确的知识。既能获得跳出 GIL 的性能优势，又能保持代码的可维护性，这证明了 Python 社区在科学计算领域付出的辛勤努力。Numba 是实现高性能和利用并行性的最佳工具之一，因此它应该成为每个 Python 开发者的工具箱。

获取软件

© . All rights reserved.