借助Python库CuPy，发掘GPU的威力

2023-11-07 51CTO

译者 | 布加迪

审校 | 重楼

CuPy简介

CuPy是一个Python/ target=_blank class=infotextkey>Python库，与NumPy和SciPy数组兼容，为GPU加速计算而设计。通过将NumPy换成CuPy语法，您可以在英伟达CUDA或AMD ROCm平台上运行代码。这让您可以使用GPU加速执行与数组相关的任务，从而更快地处理更庞大的数组。

只需换掉几行代码，就可以利用GPU的大规模并行处理能力来显著加快索引、规范化和矩阵乘法等数组操作。

CuPy还支持访问低级CUDA功能。它允许使用RawKernels将ndarray传递给现有的CUDA C/ C++程序，借助Streams简化性能，并允许直接调用CUDA Runtime API。

安装CuPy

您可以使用pip安装CuPy，但在此之前，您必须使用以下命令找到正确的CUDA版本。

!nvcc --version

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

谷歌Colab的当前版本似乎使用CUDA版本11.8。因此，我们将继续安装cupy-cuda11x版本。

如果您使用较旧的CUDA版本，我提供了下面的表格来帮助您确定要安装的合适的CuPy软件包。

借助Python库CuPy，发掘GPU的威力

选择正确的版本后，我们将使用pip安装Python软件包。

pip install cupy-cuda11x

如果您已安装了Anaconda，还可以使用conda命令自动检测并安装CuPy软件包的正确版本。

conda install -c conda-forge cupy

CuPy基础知识

在本节中，我们将比较CuPy和Numpy的语法，它们相似度为95%。您不是使用np，而是用cp代替它。

我们将先使用Python列表创建NumPy和CuPy数组。之后，我们将计算向量的范数。

import cupy as cp
import numpy as np

x = [3, 4, 5]

x_np = np.array(x)
x_cp = cp.array(x)

l2_np = np.linalg.norm(x_np)
l2_cp = cp.linalg.norm(x_cp)

print("Numpy: ", l2_np)
print("Cupy: ", l2_cp)

正如我们所见，我们得到了类似的结果。

Numpy: 7.0710678118654755

Cupy: 7.0710678118654755

为了将NumPy转换成CuPy数组，只需使用cp.asarray(X)。

x_array = np.array([10, 22, 30])
x_cp_array = cp.asarray(x_array)
type(x_cp_array)

cupy.ndarray

或者，使用.get()，将CuPy转换成Numpy数组。

x_np_array = x_cp_array.get()
type(x_np_array)

numpy.ndarray

性能比较

在本节中，我们将比较NumPy和CuPy的性能。

我们将使用time.time()对代码执行时间进行计时。然后，我们将创建一个3D NumPy数组，并执行一些数学函数。

import time

# NumPy and CPU Runtime
s = time.time()
x_cpu = np.ones((1000, 100, 1000))
np_result = np.sqrt(np.sum(x_cpu**2, axis=-1))
e = time.time()
np_time = e - s
print("Time consumed by NumPy: ", np_time)

Time consumed by NumPy: 0.5474584102630615

类似地，我们将创建一个3D CuPy数组，执行数学运算，并进行计时以评测性能。

# CuPy and GPU Runtime
s = time.time()
x_gpu = cp.ones((1000, 100, 1000))
cp_result = cp.sqrt(cp.sum(x_gpu**2, axis=-1))
e = time.time()
cp_time = e - s
print("nTime consumed by CuPy: ", cp_time)

Time consumed by CuPy: 0.001028299331665039

为了计算差异，我们将NumPy时间除以CuPy时间，我们使用CuPy后获得了500倍以上的性能提升。

diff = np_time/cp_time
print(f'nCuPy is {diff: .2f} X time faster than NumPy')

CuPy is 532.39 X time faster than NumPy

注意：为了达到更好的效果，建议进行一番预热运行，以尽量减少时间波动。

除了速度优势外，CuPy还提供了出色的多GPU支持，可以充分利用多个GPU的集体计算能力。

此外，如果您想比较结果，可以看看我的Colab笔记本。

结论

总之，CuPy提供了一种在英伟达GPU上加速NumPy代码的简单方法。只需做几个修改，将NumPy替换成CuPy，就可以体验到数组计算在速度上有数量级的提升。这种性能提升使您可以处理庞大得多的数据集和模型，从而实现更先进的机器学习和科学计算。