CuPy : NumPy & SciPy for GPU
Website | Install | Tutorial | Examples | Documentation | API Reference | Forum
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.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 |
| CUDA 11.5 | pip install cupy-cuda115 |
| ROCm 4.0 (*) | pip install cupy-rocm-4-0 |
| ROCm 4.2 (*) | pip install cupy-rocm-4-2 |
| ROCm 4.3 (*) | pip install cupy-rocm-4-3 |
(*) ROCm support is an experimental feature. Refer to the docs for details.
Use -f https://pip.cupy.dev/pre option to install pre-releases (e.g., pip install cupy-cuda114 -f https://pip.cupy.dev/pre).
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
More information
License
MIT License (see LICENSE file).
CuPy is designed based on NumPy's API and SciPy's API (see docs/LICENSE_THIRD_PARTY file).
CuPy is being maintained and developed by Preferred Networks Inc. and community contributors.
Reference
Ryosuke Okuta, Yuya Unno, Daisuke Nishino, Shohei Hido and Crissman Loomis. CuPy: A NumPy-Compatible Library for NVIDIA GPU Calculations. Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Thirty-first Annual Conference on Neural Information Processing Systems (NIPS), (2017). [PDF]
@inproceedings{cupy_learningsys2017,
author = "Okuta, Ryosuke and Unno, Yuya and Nishino, Daisuke and Hido, Shohei and Loomis, Crissman",
title = "CuPy: A NumPy-Compatible Library for NVIDIA GPU Calculations",
booktitle = "Proceedings of Workshop on Machine Learning Systems (LearningSys) in The Thirty-first Annual Conference on Neural Information Processing Systems (NIPS)",
year = "2017",
url = "http://learningsys.org/nips17/assets/papers/paper_16.pdf"
}