Numba+cuPy=GPU加速

在数值计算领域，利用 GPU 的强大功能可以显著提高性能。虽然 NumPy 长期以来一直是 Python 中数组操作和数值运算的首选库，但它缺乏原生的 GPU 加速。幸运的是，cuPy 和 Numba 这两个库通过将 GPU 支持无缝集成到您的数值工作流程中，提供了引人注目的替代方案。在这篇博文中，我们将深入研究这些库并将它们与 NumPy 进行比较，重点介绍它们的独特功能和优势。

1、cuPy

GPU 加速的 NumPy 替代品 cuPy 是一个高性能库，可模拟 NumPy API，同时提供 GPU 加速。它允许您毫不费力地将现有的 NumPy 代码转换到 GPU，为计算密集型任务带来显着的速度提升。使用 cuPy，您可以利用 NVIDIA GPU 的并行处理功能以大规模并行方式执行数组操作和数学计算。通过与 CUDA 无缝集成，cuPy 使您能够编写 GPU 加速代码，而无需进行大量修改。

示例：数组平方让我们比较一下 cuPy 与 NumPy 如何加速简单的数组平方运算：

import numpy as np
import cupy as cp

# NumPy code
a_np = np.array([1, 2, 3, 4, 5])
result_np = np.square(a_np)

# cuPy code
a_cp = cp.array([1, 2, 3, 4, 5])
result_cp = cp.square(a_cp)

# Transfer the result back to NumPy for comparison
result_cp_np = cp.asnumpy(result_cp)

print("NumPy result:", result_np)
print("cuPy result:", result_cp_np)

如示例所示，cuPy 无缝集成到 NumPy 工作流中，让您无需大量修改代码即可实现 GPU 加速。

2、Numba

GPU 的即时编译 Numba 采用不同的 GPU 加速方法。它是一个即时 (JIT) 编译器，可将 Python 函数转换为优化的机器代码，同时支持 CPU 和 GPU。通过使用 Numba 修饰函数，拟可以指示代码的哪些部分应该编译以供 GPU 执行。Numba 利用 CUDA 将计算卸载到 GPU，为拟的数值计算带来令人印象深刻的速度改进。这允许您以类似 NumPy 的方式编写代码，对其进行注释以实现 GPU 加速，并享受 GPU 并行处理的好处。

示例：矩阵乘法让我们比较一下 Numba 与 NumPy 相比如何加速矩阵乘法运算：

import numpy as np
import numba as nb

# NumPy code
a_np = np.random.rand(1000, 1000)
b_np = np.random.rand(1000, 1000)
result_np = np.matmul(a_np, b_np)

# Numba code
@nb.njit(parallel=True)
def matrix_multiply(a, b):
    return np.matmul(a, b)

result_nb = matrix_multiply(a_np, b_np)

print("NumPy result:", result_np)
print("Numba result:", result_nb)

在此示例中，Numba 的 JIT 编译和并行执行功能可显著提高速度，使其非常适合计算密集型任务。

3、比较 cuPy、Numba 和 NumPy

虽然 cuPy 和 Numba 有着 GPU 加速的共同目标，但它们提供不同的方法，并具有独特的功能，使它们彼此之间以及与 NumPy 区分开来。让我们简要比较一下这三个库：

3.1 cuPy vs. NumPy

cuPy 旨在成为 NumPy 的直接替代品，提供类似的 API 和易用性。您可以期待从 NumPy 到 cuPy 的平稳过渡，获得 GPU 加速而无需进行重大代码修改。

cuPy 允许您利用 NVIDIA GPU 的并行计算功能，从而显著加快数组操作和数学函数的计算速度。

但是，cuPy 的 GPU 加速仅限于 NVIDIA GPU，这可能会给具有不同 GPU 架构的用户带来兼容性挑战。

3.2 Numba vs. NumPy

Numba 提供 JIT 编译方法，允许您加速 CPU 和 GPU 的数值计算。这种灵活性使其适用于没有 GPU 的系统或需要跨不同硬件架构移植的系统。

Numba 需要明确注释 GPU 加速函数，这可以对将代码的哪些部分卸载到 GPU 进行细粒度控制。

另一方面，NumPy 是一个通用数值计算库，缺乏原生 GPU 支持，但被广泛采用，并拥有广泛的工具和库生态系统。

示例：使用所有三个库的图像卷积

import numpy as np
import cupy as cp
import numba as nb
from scipy.signal import convolve2d

# NumPy code
image_np = np.random.rand(512, 512)
kernel_np = np.ones((3, 3))
result_np = convolve2d(image_np, kernel_np, mode='same')

# cuPy code
image_cp = cp.asarray(image_np)
kernel_cp = cp.ones((3, 3))
result_cp = cp.convolve2d(image_cp, kernel_cp, mode='same').get()

# Numba code
@nb.njit
def image_convolution(image, kernel):
    return convolve2d(image, kernel, mode='same')

result_nb = image_convolution(image_np, kernel_np)

print("NumPy result:", result_np)
print("cuPy result:", result_cp)
print("Numba result:", result_nb)

在此示例中，所有三个库都用于使用 3x3 内核执行图像卷积。NumPy 利用 scipy.signal 中的 convolve2d 函数，cuPy 提供 GPU 加速版本的 convolve2d，Numba 使用 JIT 编译编译卷积函数。

4、结束语

cuPy 和 Numba 都为 Python 中的 GPU 加速数值计算提供了强大的 NumPy 替代方案。

cuPy 提供从 NumPy 到 GPU 加速的无缝过渡，利用 NVIDIA GPU 的并行计算能力。另一方面，Numba 为 GPU 启用 JIT 编译，允许拟编写类似 NumPy 的代码并对其进行注释以供 GPU 执行。

cuPy 和 Numba 之间的选择取决于 GPU 的可用性、所需的性能提升和兼容性要求等因素。无论选择哪个库，将 GPU 加速集成到你的数值工作流程中都可以实现显着的速度提升，并使你能够处理更苛刻的计算任务。

原文链接：Exploring GPU-Accelerated Numerical Computing: A Look into cuPy and Numba

BimAnt翻译整理，转载请标明出处