CuPy : NumPy & SciPy for GPU

CuPy is a NumPy/SciPy-compatible array library for GPU-accelerated computing with Python. CuPy acts as a drop-in replacement to run existing NumPy/SciPy code on NVIDIA CUDA or AMD ROCm platforms.

CuPy is a NumPy/SciPy-compatible array library for GPU-accelerated computing with Python.
CuPy acts as a drop-in replacement to run existing NumPy/SciPy code on NVIDIA CUDA or AMD ROCm platforms.

>>> 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 also provides access to low-level CUDA features.
You can pass ndarray to existing CUDA C/C++ programs via RawKernels, use Streams for performance, or even call CUDA Runtime APIs directly.

Installation

Wheels (precompiled binary packages) are available for Linux (x86_64) and Windows (amd64).
Choose the right package for your platform.

Platform Command
CUDA 10.0 pip install cupy-cuda100
CUDA 10.1 pip install cupy-cuda101
CUDA 10.2 pip install cupy-cuda102
CUDA 11.0 pip install cupy-cuda110
CUDA 11.1 pip install cupy-cuda111
CUDA 11.2 pip install cupy-cuda112
CUDA 11.3 pip install cupy-cuda113
CUDA 11.4 pip install cupy-cuda114
ROCm 4.0 (*) pip install cupy-rocm-4-0
ROCm 4.2 (*) pip install cupy-rocm-4-2

(*) ROCm support is an experimental feature. Refer to the docs for details.

See the Installation Guide if you are using Conda/Anaconda or building from source.

Run on Docker

Use NVIDIA Container Toolkit to run CuPy image with GPU.

$ docker run --gpus all -it cupy/cupy

GitHub

https://github.com/cupy/cupy