numba即时编译器

原创于 2025-07-19 16:54:35 发布 · 325 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能

人工智能专栏收录该内容

280 篇文章

订阅专栏

numba 是一个开源的 即时编译器（Just-In-Time, JIT），用于加速 Python 中的数值计算，尤其适用于使用 NumPy 的代码。

一、什么是 Numba？

属性	内容
名称	`numba`
类型	Python 的 JIT 编译器
支持后端	LLVM（低级虚拟机）
适用场景	数组计算、循环、数值算法
编译方式	`@jit` 或 `@njit` 装饰器
优势	代码加速，兼容性高，接近 C 的运行速度

二、为什么使用 Numba？

显著提速：Numba 编译 Python 函数为机器代码，可获得近似 C 的速度
简单易用：不需要学习 Cython 或 C++
与 NumPy 原生兼容：几乎无需改动 NumPy 代码
支持并行化：可使用多线程或 GPU 加速（通过 CUDA）

三、基本使用方法

from numba import jit
import numpy as np
import time

@jit(nopython=True)  # 推荐使用 nopython 模式获得最高性能
def sum_matrix(arr):
    total = 0.0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            total += arr[i, j]
    return total

a = np.random.rand(1000, 1000)

start = time.time()
print(sum_matrix(a))
print("耗时：", time.time() - start)

在第一次调用 sum_matrix 时，Numba 会编译它为机器代码。此后调用速度会显著提升。

四、常见装饰器说明

装饰器	说明
`@jit`	自动选择模式（但可能触发 object 模式）
`@njit`	等同于 `@jit(nopython=True)`，推荐使用
`@vectorize`	编写 ufunc 向量函数，可用于广播
`@guvectorize`	通用向量函数（支持自定义输入输出形状）
`@cuda.jit`	用于编写 GPU CUDA 核函数

五、Numba vs Cython

项目	Numba	Cython
写法	原生 Python	需写 Cython 语法或扩展
加速方式	JIT 编译	AOT 编译
上手门槛	低	较高
性能	非常高	更高（手动优化更灵活）

六、并行化和 GPU 加速

多线程加速（使用 `parallel=True`）：

from numba import njit, prange

@njit(parallel=True)
def row_sum(arr):
    result = np.zeros(arr.shape[0])
    for i in prange(arr.shape[0]):
        result[i] = np.sum(arr[i])
    return result

GPU 加速（使用 CUDA）：

from numba import cuda

@cuda.jit
def add_kernel(a, b, c):
    i = cuda.grid(1)
    if i < a.size:
        c[i] = a[i] + b[i]

# 使用 GPU

七、使用限制

不支持所有 Python 特性（如字典、类方法、异常处理）
JIT 编译有首次运行开销
不能加速解释器调用的代码（如 pandas、matplotlib）

八、典型应用场景

任务	是否适合使用 Numba
图像处理	✅ 非常适合
数值微分/积分计算	✅
科学模拟（如粒子系统）	✅
数据清洗	❌ pandas 更适合
深度学习	❌ 用 PyTorch 或 TensorFlow