Numba: 使用GPUs加速python

embedded/2024/10/24 2:54:39/

作为python函数的即时编译器，Numba是直接在标准python解释器中运行的，所以我们可以直接在GPU上运行用python写的CUDA kernels（在显卡上运行的函数），换句话说，Numba可以让我们直接使用python而非学习一门新语言的方式方便地编写在GPU上运行的CUDA kernels。

github地址：https://github.com/numba/numba

如何使用 Numba？

使用 Numba 的最常用方法是使用其装饰器集合，以此达到指示 Numba 编译指定函数的效果。当调用 Numba 修饰函数时，函数被编译为机器代码“即时” (JIT，全称是Just-In-Time) 执行，从而使指定代码以本机机器代码速度运行。[4]

使用 conda 或 pip 进行安装：

conda install numba 
# or
pip install numba

通过函数装饰器的方式使用 (以下代码示例由 ChatGPT 生成)：

python">import numpy as np
from numba import jit# 使用 numba 的 jit 装饰器加速函数
@jit(nopython=True)
def matrix_sum(A):total = 0.0for i in range(A.shape[0]):for j in range(A.shape[1]):total += A[i, j]return total# 创建一个随机的 numpy 数组
A = np.random.rand(1000, 1000)# 调用加速后的函数
result = matrix_sum(A)
print("Sum of matrix elements:", result)"""
在这个示例中：
1. @jit(nopython=True)：启用 numba 的 nopython 模式，可以确保整个函数运行时避免解释器的参与，从而加速。
2. matrix_sum 函数中使用双重循环遍历二维 numpy 数组 A，并逐个累加元素。这样可以利用 numba 的 JIT 编译来显著提高运算速度，特别是在处理大规模数据时。
"""

哪些场景下使用 Numba？

（以下内容由ChatGPT基于[5]中润色或根据提问生成）

Numba 支持 Python 的原生函数以及部分 numpy 操作，但对于更高层次的数据处理库（如 pandas），Numba 的加速效果有限，因为它无法直接优化这些库的高度封装操作 [5]。此外，Numba 也不适用于已经经过高度优化的主流机器学习框架（如 scikit-learn、TensorFlow 和 PyTorch），这些框架内部已经实现了对计算性能的优化，通常通过 C/C++ 或其他高效计算后端进行加速，因此进一步利用 Numba 进行优化不会带来显著性能提升。

如果代码大量使用 numpy 数组，Numba 可以通过 JIT 编译优化处理，避免纯 Python 循环操作，常见场景有：