Numba：一个高性能Python编译器

加速 Python 函数

Numba 使用业界标准的 LLVM 编译器库在运行时将 Python 函数翻译成优化的机器码。经 Numba 编译的 Python 数值算法可以接近 C 或 FORTRAN 的速度。

您无需替换 Python 解释器、运行单独的编译步骤，甚至不需要安装 C/C++ 编译器。只需将 Numba 的一个装饰器应用于您的 Python 函数，Numba 就会完成其余的工作。

from numba import njit
import random

@njit
def monte_carlo_pi(nsamples):
    acc = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

专为科学计算而生

Numba 旨在与 NumPy 数组和函数一起使用。Numba 为不同的数组数据类型和布局生成专门的代码以优化性能。特殊的装饰器可以创建通用函数，它们能够像 NumPy 函数一样对 NumPy 数组进行广播操作。

Numba 也能很好地与 Jupyter notebook 配合进行交互式计算，并与 Dask 和 Spark 等分布式执行框架协同工作。

了解更多 » 立即体验 »

@njit(parallel=True)
def logistic_regression(Y, X, w, iterations):
    for i in range(iterations):
        w -= np.dot(((1.0 /
              (1.0 + np.exp(-Y * np.dot(X, w)))
              - 1.0) * Y), X)
    return w

并行化您的算法

Numba 提供了一系列选项，可以并行化您在 CPU 和 GPU 上的代码，通常只需进行微小的代码改动。

简化线程

@njit(parallel=True)
def simulator(out):
    # iterate loop in parallel
    for i in prange(out.shape[0]):
        out[i] = run_sim()

Numba 可以自动在多个 CPU 核心上执行 NumPy 数组表达式，并使并行循环的编写变得容易。

了解更多 » 立即体验 »

SIMD 矢量化

LBB0_8:
    vmovups (%rax,%rdx,4), %ymm0
    vmovups (%rcx,%rdx,4), %ymm1
    vsubps  %ymm1, %ymm0, %ymm2
    vaddps  %ymm2, %ymm2, %ymm2

Numba 可以自动将某些循环翻译成矢量指令，以实现 2-4 倍的速度提升。Numba 会根据您的 CPU 功能进行调整，无论您的 CPU 支持 SSE、AVX 还是 AVX-512。

了解更多 » 立即体验 »

GPU 加速

借助于对 NVIDIA CUDA 的支持，Numba 允许您完全用 Python 编写并行 GPU 算法。

Numba CUDA »

便携式编译

交付高性能 Python 应用程序，无需担心二进制编译和打包的麻烦。您的源代码保持纯 Python，而 Numba 在运行时处理编译。我们持续在 200 多种不同的平台配置中测试 Numba。

Numba 支持 Intel 和 AMD x86、POWER8/9 以及 ARM CPU（包括 Apple M1）、NVIDIA GPU、Python 3.9-3.12，以及 Windows/macOS/Linux。大多数系统的预编译 Numba 二进制文件以 conda 包和 pip 可安装的 wheel 形式提供。

了解更多 »

致谢

Numba 的开发得益于多个组织当前和/或过去的支持。