cupy_logo

面向 Python 的 NumPy/SciPy 兼容数组库,用于 GPU 加速计算

利用 GPU 实现高性能

CuPy 是一个开源的数组库,用于使用 Python 进行 GPU 加速计算。CuPy 利用 CUDA Toolkit 库(包括 cuBLAS、cuRAND、cuSOLVER、cuSPARSE、cuFFT、cuDNN 和 NCCL)充分发挥 GPU 架构的优势。
该图显示了 CuPy 相对于 NumPy 的加速效果。大多数操作在使用 CuPy 的 GPU 上开箱即用即可良好运行。CuPy 将某些操作的速度提高了 100 多倍。请阅读 RAPIDS AI Medium 博客上的原始基准测试文章 Single-GPU CuPy Speedups

与 NumPy 和 SciPy 高度兼容

CuPy 的接口与 NumPy 和 SciPy 高度兼容;在大多数情况下,它可以作为直接替代品使用。您只需在 Python 代码中将 numpyscipy 替换为 cupycupyx.scipy 即可。CuPy 基础教程有助于学习 CuPy 的第一步。
CuPy 支持各种方法、索引、数据类型、广播等。此比较表显示了 NumPy / SciPy API 及其对应的 CuPy 实现列表。

>>> import cupy as cp
>>> x = cp.arange(6).reshape(2, 3).astype('f')
>>> x
array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.]], dtype=float32)
>>> x.sum(axis=1)
array([  3.,  12.], dtype=float32)

易于安装

安装 CuPy 最简单的方法是使用 pip。CuPy 为 Linux 和 Windows 提供了 wheels(预编译的二进制包)。请阅读安装指南了解更多详情。
CuPy 也可以从 Conda-Forge 安装从源代码安装

# For CUDA 11.2 ~ 11.x
pip install cupy-cuda11x

# For CUDA 12.x
pip install cupy-cuda12x

# For AMD ROCm 4.3
pip install cupy-rocm-4-3

# For AMD ROCm 5.0
pip install cupy-rocm-5-0

易于编写自定义内核

如果您想让代码运行得更快,您可以轻松创建自定义 CUDA 内核,只需一小段 C++ 代码片段即可。CuPy 会自动封装并编译它以生成 CUDA 二进制文件。编译后的二进制文件会被缓存并在后续运行中重用。请阅读用户定义内核教程。
此外,您还可以通过原始模块 (Raw modules)使用原始 CUDA 内核。

>>> x = cp.arange(6, dtype='f').reshape(2, 3)
>>> y = cp.arange(3, dtype='f')
>>> kernel = cp.ElementwiseKernel(
...     'float32 x, float32 y', 'float32 z',
...     '''
...     if (x - 2 > y) {
...       z = x * y;
...     } else {
...       z = x + y;
...     }
...     ''', 'my_kernel')
>>> kernel(x, y)
array([[ 0.,  2.,  4.],
       [ 0.,  4.,  10.]], dtype=float32)

观看视频。

PyBay 2019,旧金山
SciPy Japan Conference 2020,线上

开发者


支持方