Numba 是一个开源的 JIT(Just-In-Time)编译器,用于加速 Python 代码。它可以将 Python 函数转换为快速的机器码,从而提高代码的执行速度。要使用 Numba 加速 Python 代码,请按照以下步骤操作:
-
安装 Numba: 你可以使用 pip 安装 Numba。打开命令行或终端,然后输入以下命令:
pip install numba
-
使用
@jit
装饰器: 在你想要加速的函数上方添加@jit
装饰器。这将使用 Numba 的 JIT 编译器对该函数进行优化。例如:import numba @numba.jit(nopython=True) def my_function(a, b): return a + b
这里,
nopython=True
表示我们希望 Numba 尽最大努力不使用 Python 的动态类型特性,以便生成更快的代码。 -
使用
numba.jitclass
装饰器: 对于包含多个方法的类,你可以使用numba.jitclass
装饰器将其转换为 Numba 的 JIT 类。例如:import numba @numba.jitclass({ 'a': numba.int32, 'b': numba.int32, 'add': numba.int32(numba.int32, numba.int32) }) class MyClass: def __init__(self, a, b): self.a = a self.b = b def add(self): return self.a + self.b
然后,你可以像普通类一样实例化并使用这个类:
obj = MyClass(1, 2) result = obj.add() print(result) # 输出 3
-
使用
numba.vectorize
装饰器: 如果你有一个接受单个数组参数的函数,并希望将其加速为数组操作,可以使用numba.vectorize
装饰器。例如:import numba @numba.vectorize([numba.float64(numba.float64, numba.float64)]) def add_arrays(a, b): return a + b x = [1.0, 2.0, 3.0] y = [4.0, 5.0, 6.0] z = add_arrays(x, y) print(z) # 输出 [5.0, 7.0, 9.0]
-
使用
numba.cuda.jit
装饰器: 如果你想在 GPU 上加速你的代码,可以使用numba.cuda.jit
装饰器。这将生成 CUDA 核函数,可以在 NVIDIA GPU 上执行。例如:import numba from numba import cuda @cuda.jit def add_arrays_gpu(a, b, c): i = cuda.grid(1) if i < len(a): c[i] = a[i] + b[i] x = [1, 2, 3] y = [4, 5, 6] z = [0] * len(x) add_arrays_gpu[(1, 1)](x, y, z) print(z) # 输出 [5, 7, 9]
通过以上方法,你可以使用 Numba 加速 Python 代码。请注意,Numba 对某些类型的代码(如包含循环和递归的代码)的加速效果可能不佳。在这种情况下,你可能需要寻找其他优化方法。