Numba加速计算（CPU + GPU + prange）

加速方法：Numba、CuPy、PyTorch、PyCUDA、Dask、Rapids

加速方法	简介	支持平台	适用范围
Numba	通过即时编译（JIT）来加速Python函数	CPU + GPU	数值计算（密集循环）
CuPy	NumPy的GPU加速库	GPU	数组操作和数学计算（大规模数据集）
PyTorch	深度学习框架	GPU	张量操作和自动求导
PyCUDA	与CUDA交互的Python库	GPU	在Python中编写CUDA代码并在GPU上执行
Dask	并行计算库	CPU + GPU	并行操作（大规模数据集）
Rapids	基于NVIDIA GPU加速的数据科学生态系统	GPU	提供数据处理和机器学习的库（如cuDF、cuML）

一、Numba简介

Numba官网：专为 NumPy 科学计算而打造的，用于加速 Python 代码的即时编译器（Just-In-Time, JIT Compiler）。

由 Anaconda 公司主导开发

• 原理：使用行业标准LLVM编译器库在运行时，将 Python 和 NumPy 代码的子集转换为快速的机器代码。
• 速度：接近 C 或 FORTRAN 的速度
• 操作简单：只需将 Numba 装饰器之一应用到 Python 函数，Numba 将使用即时编译（JIT）编译为原生代码（机器代码），然后加速运行。
  • 不需要替换 Python 解释器、运行单独的编译步骤，也不需要安装 C/C++ 编译器）
  • 原生代码（Native code）：直接在计算机硬件上执行的机器代码。

适用范围

• 只支持NumPy库（Numba 基于 NumPy 底层代码开发）
• 不支持其余的Python库（自定义-重开发：将函数分解为底层代码）
  • （1）NumPy数值计算（CPU）：将 Python 函数即时编译为机器代码，用于加速数学运算等计算密集型任务（最初设计）。
  • （2）NumPy数组操作（CPU）：高效处理 NumPy 的数组操作和广播操作，且可以加速大型数组操作。
  • （3）支持并行计算（CPU和GPU）：使用 prange 来并行处理循环结构
  • （4）支持GPU加速：将代码转移到 GPU 上以加速执行

二、Numba类型：CPU + GPU

Numba官网案例1：并行化测试（@jit、prange、dask）
Numba官网案例2：GPU加速

import numba  # pip install numba
from numba import cuda, jit##############################################################
"""Numba装饰器：CPU加速"""
函数说明：@jit(nopython=True, parallel=True, target='cpu')
输入参数：（1）加速模式：nopython=True（默认）、forceobj=True			备注：Numba默认使用nopython编译函数。若无法完成编译，将使用对象模式（将导致性能损失）。（2）并行模式：parallel=True（默认）（3）指定平台：target='cpu'（默认）、target='gpu'> 多种不同的应用@jit()				# 适用于科学计算、数值计算和密集计算。@jit(target='cpu')	# CPU加速（与@jit()等效）@jit(target='cuda')	# GPU加速（与@cuda.jit()等效）@jit(nopython=True)	# 强制使用 nopython 模式，将函数尽可能编译成机器代码，如果无法完成编译则会引发错误（与@njit()等效）。@jit(forceobj=True)	# 强制使用对象模式，而不是默认的 nopython 模式。适用于一些特殊情况，如涉及动态类型的代码。@jit(parallel=True)	# 尝试并行化循环，充分利用多核处理器的性能（需将range转换为prange）。@njit()			# 强制使用 nopython 模式，比@jit更快但更严格（只接受Numpy数据类型）。@vectorize		# 适用于元素级别的向量化操作（单输入和单输出）、（输出数组的形状由输入数组的形状决定）@guvectorize	# 适用于元素级别的向量化操作（多输入和输出数组）、（输出数组的形状可以指定）@stencil		# 一种基于固定模板的局部计算。通过访问输入数组的邻域元素来计算输出数组的每个元素。
##############################################################
"""Numba装饰器：GPU加速"""	@cuda.jit()		# 用于在GPU上执行 CUDA 加速# @cuda.jit()与 @cuda.jit ———— 若不传递参数，两者是等效的。
# 		@cuda.jit 	使用默认选项 @cuda.jit()，并且不需要传递任何参数。
# 		@cuda.jit() 显式调用 @cuda.jit 装饰器，并且可以传递一些选项参数。

三、项目实战 —— 数组的每个元素加2

3.1、使用 python - range 循环计算 —— （时耗：137.37 秒）

import numpy as npdef numpy_cpu_kernel(input_array):shape = input_array.shape  # 获取数组形状result_array = np.zeros(shape)  # 初始化一个全零数组，形状与输入数组相同# 循环遍历每个元素，将其加 2for z in range(shape[0]):for y in range(shape[1]):for x in range(shape[2]):result_array[z, y, x] = input_array[z, y, x] + 2return result_arrayif __name__ == "__main__":# 在主机上创建 3D 数组input_data = np.zeros((1024, 1024, int(1024 * 0.5)))import timestart_time = time.time()# 在主机上调用 NumPy 函数result_array_on_host = numpy_cpu_kernel(input_data)print(f"总共耗时: {time.time() - start_time:.2f} 秒")# 打印结果数组的形状和最大值print(result_array_on_host.shape)print(result_array_on_host.max())"""总共耗时: 137.37 秒"""

3.2、使用 python - numpy 数组计算 —— （时耗：1.05 秒）

import numpy as npdef numpy_cpu_kernel(input_array):return input_array + 2  # 将输入数组的每个元素加 2if __name__ == "__main__":# 在主机上创建 3D 数组input_data = np.zeros((1024, 1024, int(1024 * 0.5)))import timestart_time = time.time()# 在主机上调用 NumPy 函数result_array_on_host = numpy_cpu_kernel(input_data)print(f"总共耗时: {time.time() - start_time:.2f} 秒")# 打印结果数组的形状和最大值print(result_array_on_host.shape)print(result_array_on_host.max())"""总共耗时: 1.05 秒"""

3.3、使用 numba - CPU 加速计算 —— （时耗：13.85 秒）

from numba import jit
import numpy as np# 使用 Numba 的 jit 装饰器进行即时编译
@jit(nopython=True)
def my_cpu_kernel(input_array, output_array):# 使用三个嵌套循环遍历 3D 输入数组的每个元素for x in range(input_array.shape[2]):for z in range(input_array.shape[0]):for y in range(input_array.shape[1]):# 执行简单的操作：将输入数组的元素加 2，并将结果存储到输出数组中output_array[z, y, x] = input_array[z, y, x] + 2if __name__ == "__main__":# （1）在主机上创建输入数组input_data = np.zeros((1024, 1024, int(1024 * 0.5)))  # 创建一个全零的3D数组result_array_on_host = np.zeros_like(input_data)# （2）在 CPU 上调用加速函数import timestart_time = time.time()my_cpu_kernel(input_data, result_array_on_host)print(f"总运行时间: {time.time() - start_time:.2f} 秒")  # 打印运行时间# 打印输出数组的形状和最大值print(result_array_on_host.shape)print(result_array_on_host.max())"""总共耗时: 13.858259439468384"""

3.4、使用 numba - GPU 加速计算 —— （时耗：0.13 秒）

【深度学习环境配置】Anaconda +Pycharm + CUDA +cuDNN + Pytorch + Opencv（资源已上传）

在CUDA编程中

• CPU和主存（RAM）：称为主机（Host）
• GPU和显存（VRAM）：称为设备（Device）
  • CPU无法直接读取显存数据，GPU无法直接读取主存数据；
  • 主机与设备必须通过总线（Bus）相互通讯；

RAM是CPU的主内存，显存是GPU的专用内存

GPU计算流程：

• （1）cuda.to_device()：将主机端的数据拷贝到设备端上，并在GPU上分配与主机上数据相同大小的内存。
• （2）cuda.device_array_like()或cuda.device_array()：在GPU上分配用于输出数据的内存。
• （3）gpu[blocks_per_grid, threads_per_block]：在CPU上调用GPU函数，启动GPU多核并行计算（详细看1.2）。
• （4）CPU与GPU异步计算；
  • GPU函数的启动方式是异步的
    • 异步计算：CPU不会等待GPU函数执行完毕才执行下一行代码。
    • 同步计算：在调用的GPU函数后面添加 cuda.synchronize() —— 表示CPU需要等待GPU函数执行后再计算。
• （5）cuda.copy_to_host()：将GPU设备端的计算结果拷贝回CPU主机端上。Python通过Numba实现GPU加速

1. numba的GPU加速 —— 1天到1分钟的转变
2. numba的GPU加速 —— 超过Numpy的速度有多难？

from numba import cuda
import numpy as np# 使用 Numba 的 CUDA 装饰器进行 GPU 加速
@cuda.jit
def my_cuda_kernel(input_array, output_array):x = cuda.grid(1)  # 使用1维索引if x < input_array.shape[2]:  # 检查索引是否在数组范围内for z in range(input_array.shape[0]):for y in range(input_array.shape[1]):# 执行简单的操作：将输入数组的元素加 2，并将结果存储到输出数组中output_array[z, y, x] = input_array[z, y, x] + 2if __name__ == "__main__":# （1）在主机上创建3D数组input_data = np.zeros((1024, 1024, int(1024 * 0.5)))# （2）拷贝数据 + 在GPU上分配内存device_input_array = cuda.to_device(input_data)  # 将主机上的input_data复制到GPU上device_output_array = cuda.device_array_like(input_data)  # 在GPU上开辟一个与 input_data 相同形状的数组用于存储计算结果。# （3）定义线程块的大小 + 线程块的数量threads_per_block = (16,)blocks_per_grid_x = (input_data.shape[2] + threads_per_block[0] - 1) // threads_per_block[0]blocks_per_grid = (blocks_per_grid_x,)# （4）调用CUDA核函数import timestart_time = time.time()my_cuda_kernel[blocks_per_grid, threads_per_block](device_input_array, device_output_array)print(f"总运行时间: {time.time() - start_time:.2f} 秒")  # 打印运行时间# （5）将结果从GPU复制回主机result_array_on_host = device_output_array.copy_to_host()  # 将计算结果从GPU复制回主机print(result_array_on_host.shape)print(result_array_on_host.max())"""总共耗时: 0.136000394821167"""

3.5、使用 numba.prange 并行循环计算

prange（parallel range）：类似于Python标准库中的 range，但专用于并行化循环计算（引入额外的开销）。可以在多个处理器核心上同时执行循环，以提高密集型计算的性能。

适用范围：

• 大规模数据和密集型计算：只有在循环非常庞大时，才可以充分利用多核处理器，否则将导致耗时更长。
• 独立迭代：每个迭代之间不存在依赖关系时，若后面的迭代依赖于前面迭代的计算结果（类似于递归函数），可能会导致耗时更长。
• NumPy数组：使用NumPy数组可以表现出更好的性能。

import random
import numba
import time###############################################
# numba加速 + 串行版本（Serial version）
###############################################
@numba.jit(nopython=True)
def monte_carlo_pi_serial(nsamples):acc = 0for i in range(nsamples):x = random.random()y = random.random()if (x ** 2 + y ** 2) < 1.0:acc += 1return 4.0 * acc / nsamples###############################################
# numba加速 + 并行版本（Parallel version）
###############################################
@numba.jit(nopython=True, parallel=True)
def monte_carlo_pi_parallel(nsamples):acc = 0for i in numba.prange(nsamples):x = random.random()y = random.random()if (x ** 2 + y ** 2) < 1.0:acc += 1return 4.0 * acc / nsamplesif __name__ == "__main__":start_time = time.time()monte_carlo_pi_serial(33000000)  # 串行版本print(f"总共耗时: {time.time() - start_time:.2f} 秒")start_time = time.time()monte_carlo_pi_parallel(33000000)  # 并行版本print(f"总共耗时: {time.time() - start_time:.2f} 秒")"""一万次循环:0.277700185775756840.42730212211608887三千万次循环:0.464895009994506840.4749984741210937510亿次循环:6.72973990440368651.1980292797088623100亿次循环:67.732758283615116.892062425613403"""