Numba基础

embedded/2024/9/25 22:30:54/

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 = 0for i in arr:  # 遍历数组元素并累加total += ireturn 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 = 1for i in arr:  # 遍历数组元素并累乘total *= ireturn 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}")

解释:

  • prangerange 的并行版本,它将循环拆分为多个线程并行执行,从而充分利用多核 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 加速功能来进一步提升性能。


http://www.ppmy.cn/embedded/113218.html

相关文章

k8s的搭建

一、安装环境 准备三台主机&#xff1a; 192.168.1.66 k8s-master 192.168.1.77 k8s-node01 192.168.1.88 k8s-node02 网段&#xff1a; Pod ⽹段 172.16.0.0/16 Service ⽹段 10.96.0.0/16 注&#xff1a;宿主机⽹段、Pod…

VMware Tools系列二:图解安装VMware Tools过程

一、安装环境&#xff1a; VMware Workstation Pro17华为OpenEuler虚拟机 二、安装步骤&#xff1a; 在VMware Workstation中为OpenEuler虚拟机安装VMware Tools的步骤如下&#xff1a; 1. 确保OpenEuler虚拟机正在运行。注意&#xff0c;安装的VMware Tools只对当前虚拟机…

『功能项目』切换职业技能面板【49】

我们打开上一篇48切换职业面板的项目&#xff0c; 本章要做的事情是制作第二职业法师技能面板、第三职业面板并且完成切换 双击打开Canvas进入预制体空间 复制三个技能栏面板 重命名 设置第一技能栏 设置第二职业技能栏 设置第三职业技能栏 修改脚本&#xff1a;ChangeProfess…

雷达探测系统中的显示模式详解: P显、B显、E显、A显、H显

雷达探测系统中的显示模式详解: P显、B显、E显、A显、H显 在雷达探测系统中&#xff0c;显示模式是雷达设备的重要组成部分&#xff0c;不同的显示模式可以提供不同维度的目标信息&#xff0c;以满足多种应用需求。常见的雷达显示模式包括 P 显、B 显、E 显、A 显 和 H 显。本…

共享内存C(Linux)

在学习的时候遇到问题&#xff0c;就是将结构体作为共享内存时将string类型置入结构体内&#xff0c;导致程序出现段错误&#xff0c;后来经过排查发现共享内存是c语言的库不支持string类型&#xff0c;需要用char name[20]代替。 1.在Linux中如何查看共享内存 &#xff08;1&…

93、k8s之hpa+helm

一、HPA HPA: pod的数量进行扩缩容 针对对控制器创建的pod deployment&#xff1a; replicas: 静态&#xff1a;edit yaml&#xff1a;apply -f HPA&#xff1a;基于cpu的利用率来实现pod的数量大的自动伸缩。 Horizontal Pod Autoscaling yaml 文件------主流------…

Unity persistentDataPath使用案例

Unity persistentDataPath使用案例 一、Application.persistentDataPath 1、概念 persistentDataPath&#xff1a;此属性用于返回一个持久化数据存储目录的路径&#xff0c;可以在此路径下存储一些持久化的数据文件&#xff1b;是一个可读写的目录&#xff1b;此文件夹在Edi…

如何优雅地处理返回值

我们已经知道了如何优雅的校验传入的参数了&#xff0c;那么后端服务器如何实现把数据返回给前端呢&#xff1f; 返回格式 后端返回给前端我们一般用 JSON 体方式&#xff0c;定义如下&#xff1a; {#返回状态码code:string, #返回信息描述message:string,#返回值data…