Numba基础
1. Numba 基础
1.1 什么是 Numba？
Numba 是一个 JIT 编译器，用于加速数值计算。它通过即时编译技术，将 Python 代码在运行时编译为机器代码，极大地提升执行速度，特别适合循环和矩阵操作等密集型计算。
2. Numba 基本使用
2.1 @jit 装饰器
@jit 装饰器是 Numba 的基本加速手段，它将被装饰的函数动态编译为机器码，以提高性能。如果设置 nopython=True，它会进入无 Python 模式，完全避免 Python 的解释开销。
示例代码：
import numpy as np
from numba import jit
# 使用 @jit 加速计算，启用 nopython 模式，保证最高的性能
@jit(nopython=True)
def sum_array(arr):
total = 0
for i in arr: # 遍历数组元素并累加
total += i
return total
# 生成一个随机的大型数组作为输入数据
data = np.random.rand(1000000)
# 执行加速后的函数并输出结果
result = sum_array(data)
print(f"Sum result: {result}")
解释：
@jit(nopython=True)指示 Numba 进入无 Python 模式，避免 Python 解释器的参与，获得最大性能提升。sum_array是一个简单的数组求和函数，由于使用了循环操作，因此 Numba 能够大幅加速执行。
2.2 @njit 装饰器
@njit 是 @jit(nopython=True) 的简写，功能相同，确保 Numba 进入无 Python 模式。
from numba import njit
# 使用 @njit 直接进入无 Python 模式
@njit
def multiply_array(arr):
total = 1
for i in arr: # 遍历数组元素并累乘
total *= i
return total
# 生成随机数组并计算其元素的乘积
data = np.random.rand(1000000)
result = multiply_array(data)
print(f"Product result: {result}")
解释：
@njit是一个简便的写法，相当于@jit(nopython=True)，用于自动加速函数。- 这里的
multiply_array函数执行数组的累乘操作，Numba 对这种循环密集型任务有显著的加速效果。
3. Numba 与 NumPy
Numba 对 NumPy 的支持非常友好，它能够识别并加速 NumPy 的许多函数，特别是在处理大型矩阵和数组时，可以极大提升性能。
示例代码：使用 NumPy 和 Numba 加速矩阵运算
import numpy as np
from numba import njit
# 使用 Numba 加速矩阵乘法
@njit
def matrix_mult(A, B):
# 初始化结果矩阵为零矩阵
C = np.zeros((A.shape[0], B.shape[1]))
# 三重循环执行矩阵乘法
for i in range(A.shape[0]):
for j in range(B.shape[1]):
for k in range(A.shape[1]):
# 将A的行与B的列相乘累加
C[i, j] += A[i, k] * B[k, j]
return C
# 生成两个500x500的随机矩阵
A = np.random.rand(500, 500)
B = np.random.rand(500, 500)
# 执行加速后的矩阵乘法
C = matrix_mult(A, B)
print(C)
解释：
- 这里的
matrix_mult是一个经典的三重循环矩阵乘法实现。由于是循环密集型操作，Numba 能够对其进行有效加速。 - 函数内使用了 NumPy 的
np.zeros来初始化结果矩阵，这也是被 Numba 支持并优化的 NumPy 操作。
4. Numba 并行化
Numba 提供了并行化支持，允许在多核 CPU 上同时执行任务，提升性能。通过 @njit(parallel=True) 和 prange，你可以轻松并行化代码中的循环。
4.1 使用 prange 并行化
prange 是并行版本的 range，可以将循环的不同部分分配到多个线程中执行。
示例代码：并行化求和
from numba import njit, prange
import numpy as np
# 使用并行化加速求和
@njit(parallel=True)
def parallel_sum(arr):
total = 0
# 使用 prange 代替 range 实现并行化循环
for i in prange(len(arr)):
total += arr[i] # 各线程并行计算不同部分的数组求和
return total
# 生成随机数组
data = np.random.rand(1000000)
# 执行并行求和
result = parallel_sum(data)
print(f"Parallel Sum Result: {result}")
解释：
prange是range的并行版本，它将循环拆分为多个线程并行执行，从而充分利用多核 CPU。@njit(parallel=True)告诉 Numba 对这个函数进行并行优化。
4.2 并行化向量运算
对于一些简单的向量操作，比如数组归一化，Numba 的并行化也能提供很好的加速。
示例代码：并行化的数组归一化
@njit(parallel=True)
def normalize(arr):
n = len(arr)
result = np.empty(n) # 初始化结果数组
total = 0
# 第一次并行循环计算数组的总和
for i in prange(n):
total += arr[i]
mean = total / n # 计算平均值
# 第二次并行循环进行归一化操作
for i in prange(n):
result[i] = arr[i] / mean # 将每个元素除以均值
return result
# 生成随机数组并进行归一化
data = np.random.rand(1000000)
normalized_data = normalize(data)
print(normalized_data[:10]) # 打印归一化结果的前10个元素
解释：
- 在这个例子中，
prange用于对两个独立的循环并行化：一个用于求和，另一个用于归一化。 - Numba 可以高效地并行化向量操作，提高处理大规模数组的效率。
5. GPU 加速
Numba 支持使用 CUDA 将计算任务卸载到 GPU 上执行。GPU 非常擅长处理大规模并行计算，尤其是在矩阵运算和深度学习等领域。
5.1 安装 CUDA 支持
要使用 Numba 的 GPU 功能，首先你需要安装 NVIDIA CUDA Toolkit，并确保有一个支持 CUDA 的 NVIDIA GPU。
5.2 使用 CUDA 加速
示例代码：简单的 GPU 向量加法
from numba import cuda
import numpy as np
# 定义一个在 GPU 上运行的内核函数
@cuda.jit
def gpu_add(a, b, c):
i = cuda.grid(1) # 获取线程的唯一ID（索引）
if i < c.size: # 确保索引在数组范围内
c[i] = a[i] + b[i] # 执行数组加法
# 初始化输入数据
n = 1000000
a = np.random.rand(n)
b = np.random.rand(n)
c = np.zeros(n)
# 将数据复制到 GPU
a_device = cuda.to_device(a)
b_device = cuda.to_device(b)
c_device = cuda.to_device(c)
# 设置 GPU 线程数和块数
threads_per_block = 512
blocks_per_grid = (a_device.size + (threads_per_block - 1)) // threads_per_block
# 调用 GPU 上的加法函数
gpu_add[blocks_per_grid, threads_per_block](a_device, b_device, c_device)
# 将结果从 GPU 复制回 CPU
c = c_device.copy_to_host()
print(c[:10]) # 打印结果的前10个元素
解释：
@cuda.jit定义了一个 CUDA 内核函数，在 GPU 上并行执行数组的元素加法。cuda.grid(1)获取当前线程的索引，以确定每个线程处理的数据块。- 数据通过
cuda.to_device传输到 GPU，GPU 完成计算后，再通过copy_to_host将结果返回到 CPU。
总结
Numba 是一个强大的工具，它能够极大地加速 Python 代码，尤其是在数值计算和矩阵操作方面。通过 @jit 或 @njit，可以轻松地将 Python 代码编译为高效的机器代码，同时可以利用并行化和 GPU 加速功能来进一步提升性能。