numba

Numba - 使用 JIT 加速 Python

Numba 能让 Python 代码运行得更快。它通过使用装饰器修饰函数来告诉 Numba 编译这些函数。它尤其适用于包含大量数值循环和 NumPy 数组操作的代码。

适用场景

• NumPy 内置的向量化无法满足特定算法的需求。
• 存在复杂的嵌套循环，在标准 Python 中运行缓慢。
• 需要编写逐元素处理数组的自定义“ufunc”（通用函数）。
• 高性能物理模拟（蒙特卡洛、N 体模拟、基于网格的求解器）。
• 加速代码在 NVIDIA GPU（CUDA）上执行。
• 创建可利用所有 CPU 核心的并行代码，且无需承担多进程的开销。

参考文档

官方文档：https://numba.pydata.org/numba-doc/latest/index.html
用户指南：https://numba.pydata.org/numba-doc/latest/user/index.html
搜索模式：@njit、@vectorize、prange、cuda.jit、numba.typed

核心原则

nopython 模式（@njit）

这是 Numba 的“黄金标准”。在此模式下，Numba 无需使用 Python C-API 即可编译代码，从而实现最高速度。如果无法编译（例如，因为使用了不受支持的 Python 对象），则会抛出错误。

即时（JIT）编译

函数在首次调用时进行编译。之后的调用会使用缓存的机器码。

面向数组

Numba 专为处理 NumPy 数组而设计。它能理解数组的内存布局，并可为数组生成高度优化的循环。

快速参考

安装

pip install numba

标准导入

import numpy as np
from numba import njit, prange, vectorize, guvectorize, cuda

基本模式 - 加速循环

import numpy as np
from numba import njit

# 1. Apply the @njit decorator (alias for @jit(nopython=True))
@njit
def sum_array(arr):
    res = 0.0
    # Standard Python loop that would be slow is now fast as C
    for i in range(arr.shape[0]):
        res += arr[i]
    return res

# 2. Execute
data = np.random.random(1_000_000)
result = sum_array(data) # First call compiles, then runs

重要规则

✅ 应该做

• 优先使用 @njit - 始终使用 nopython=True（或其别名 @njit）。这能确保代码确实以机器速度运行。
• 使用 NumPy 数组 - Numba 针对 NumPy 进行了优化。在 JIT 编译的函数中避免使用标准 Python 列表。
• 启用并行 - 使用 @njit(parallel=True)，并用 prange 替代 range，以实现自动多线程。
• 缓存编译后的代码 - 使用 @njit(cache=True)，避免每次重启脚本时都重新编译。
• 预热 - 请记住，首次调用会因编译而较慢。进行计时基准测试时，务必先运行一次函数再测量。
• 指定类型（可选） - 可以提供签名（例如 (float64[:],)）来加快首次调用，但 Numba 通常能很好地推断类型。

❌ 不要做

• 不要使用 Python 对象 - 在 nopython 模式下，字符串、字典和自定义类速度较慢或不受支持。如有需要，可使用 numba.typed 创建专用容器。
• 不要对小函数进行 JIT 编译 - 从 Python 调用 JIT 编译函数的开销，可能超过简单操作带来的收益。
• 不要使用不受支持的库 - 不能在 @njit 函数中使用 pandas、matplotlib 或 requests。
• 不要修改全局状态 - 为了保持稳定，JIT 编译的函数应尽可能保持“纯函数”特性。

反模式（绝对不要）

from numba import njit
import pandas as pd

# ❌ BAD: Using Pandas inside @njit (Unsupported)
@njit
def bad_func(df):
    return df['col'].sum() # Will raise a LoweringError

# ✅ GOOD: Pass NumPy arrays instead
@njit
def good_func(arr):
    return arr.sum()

# ❌ BAD: Using @jit without nopython=True
from numba import jit
@jit
def slow_func(x): # This might fall back to "Object Mode" (slow)
    return x + 1

# ✅ GOOD: Always ensure nopython mode
@njit
def fast_func(x):
    return x + 1

# ❌ BAD: Manual loops in Python to call a JIT function
# for i in range(1000):
#     process_element(arr[i]) # Calling JIT overhead 1000 times

# ✅ GOOD: Move the loop INSIDE the @njit function
@njit
def process_all(arr):
    for i in range(arr.shape[0]):
        process_element(arr[i])

并行与向量化

自动多线程

from numba import njit, prange

@njit(parallel=True)
def parallel_sum(A):
    # Use prange for the loop that should be parallelized
    s = 0
    for i in prange(A.shape[0]):
        s += A[i]
    return s

创建快速 ufunc（@vectorize）

from numba import vectorize

# This creates a NumPy ufunc that supports broadcasting
@vectorize(['float64(float64, float64)'], target='parallel')
def fast_add(x, y):
    return x + y

# Now you can use it on massive arrays
res = fast_add(arr1, arr2)

使用结构体和类型

使用 numba.typed 处理非数组数据

from numba.typed import List, Dict
from numba import njit

@njit
def use_typed_list():
    l = List()
    l.append(1.0)
    return l

GPU 加速（numba.cuda）

编写 CUDA 内核

from numba import cuda

@cuda.jit
def my_kernel(io_array):
    # Calculate thread indices
    pos = cuda.grid(1)
    if pos < io_array.size:
        io_array[pos] *= 2

# Usage
data = np.ones(256)
threadsperblock = 32
blockspergrid = (data.size + (threadsperblock - 1)) // threadsperblock
my_kernel[blockspergrid, threadsperblock](data)

实用工作流

1. 快速蒙特卡洛模拟

import random

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

2. 自定义图像滤镜（模板）

from numba import njit

@njit
def apply_threshold(image, threshold):
    M, N = image.shape
    result = np.zeros_like(image)
    for i in range(M):
        for j in range(N):
            if image[i, j] > threshold:
                result[i, j] = 255
    return result

3. 求解物理网格（拉普拉斯方程）

@njit
def solve_laplace(u, niters):
    M, N = u.shape
    for n in range(niters):
        for i in range(1, M-1):
            for j in range(1, N-1):
                u[i, j] = 0.25 * (u[i+1, j] + u[i-1, j] + u[i, j+1] + u[i, j-1])
    return u

性能优化

inspect_types() 方法

使用此方法可查看 Numba 是否不得不回退到开销较大的 Python 对象，或者是否成功将所有内容优化为原生类型。

fast_func.inspect_types() # Prints color-coded annotated code

避免在循环中分配数组

在 @njit 函数外预先分配数组，或将其作为参数传入，以避免内存管理开销。

# ✅ GOOD:
@njit
def compute_into(out_arr, in_arr):
    for i in range(in_arr.shape[0]):
        out_arr[i] = in_arr[i] * 2

常见陷阱与解决方案

“全局变量”问题

Numba 会在编译时捕获全局变量的值。

# ❌ Problem: Changing a global variable won't affect the jitted function
K = 10
@njit
def f(x): return x + K

K = 20
f(1) # Result is still 11!

# ✅ Solution: Pass constants as arguments

对象模式回退

如果 Numba 提示“已启用对象模式”，说明代码运行速度会很慢。

# ✅ Solution: Force nopython mode
@njit # If this throws error, fix the code instead of removing @njit

并行模式下的随机种子

在 parallel=True 中使用 np.random 时，需谨慎处理，以确保每个线程使用独立的随机数流。在 Numba 中使用标准的 random.random() 或 np.random.random() 是线程安全的，并会自动为每个线程处理种子。

最佳实践

1. 始终使用 @njit——不要在未设置 nopython=True 的情况下使用 @jit
2. 预先分配数组——避免在热点循环中创建数组
3. 使用 prange 实现并行——通过 parallel=True 和 prange 启用自动多线程
4. 缓存已编译函数——使用 cache=True 避免重复编译
5. 预热函数——基准测试前先调用一次 JIT 编译函数
6. 传递 NumPy 数组——调用 JIT 编译函数前，将 Python 列表转换为 NumPy 数组
7. 避免使用 Python 对象——需要容器时，使用 numba.typed.List 和 numba.typed.Dict
8. 检查编译模式——使用 inspect_types() 验证是否处于 nopython 模式
9. 处理首次调用开销——请记住，首次调用会编译函数
10. 使用合适的签名——此项为可选，但可以加快首次编译

Numba 是连接 Python 与 C++、Fortran 高性能计算领域的桥梁。它消除了循环中的“Python 开销”，让你无需牺牲执行速度即可快速进行原型开发。