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

文章目录

加速方法：Numba、CuPy、PyTorch、PyCUDA、Dask、Rapids
一、Numba简介
二、Numba类型：CPU + GPU
三、项目实战 —— 数组的每个元素加2
- 3.1、使用 python - range 循环计算 —— （时耗：137.37 秒）
- 3.2、使用 python - numpy 数组计算 —— （时耗：1.05 秒）
- 3.3、使用 numba - CPU 加速计算 —— （时耗：13.85 秒）
- 3.4、使用 numba - GPU 加速计算 —— （时耗：0.13 秒）
- 3.5、使用 numba.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 np

def numpy_cpu_kernel(input_array):
    shape = input_array.shape  # 获取数组形状
    result_array = np.zeros(shape)  # 初始化一个全零数组，形状与输入数组相同
    # 循环遍历每个元素，将其加 2
    for 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] + 2
    return result_array

if __name__ == "__main__":
    # 在主机上创建 3D 数组
    input_data = np.zeros((1024, 1024, int(1024 * 0.5)))

    import time
    start_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 np

def numpy_cpu_kernel(input_array):
    return input_array + 2  # 将输入数组的每个元素加 2

if __name__ == "__main__":
    # 在主机上创建 3D 数组
    input_data = np.zeros((1024, 1024, int(1024 * 0.5)))

    import time
    start_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] + 2


if __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 time
    start_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加速

numba的GPU加速 —— 1天到1分钟的转变
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] + 2

if __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 time
    start_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 = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

###############################################
# numba加速 + 并行版本（Parallel version）
###############################################
@numba.jit(nopython=True, parallel=True)
def monte_carlo_pi_parallel(nsamples):
    acc = 0
    for i in numba.prange(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples


if __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.27770018577575684
        0.42730212211608887

    三千万次循环:
        0.46489500999450684
        0.47499847412109375

    10亿次循环:
        6.7297399044036865
        1.1980292797088623

    100亿次循环:
        67.73275828361511
        6.892062425613403
    """