CuPy

 

CuPy is an open-source matrix library accelerated with NVIDIA CUDA.

CuPy provides GPU accelerated computing with Python.

CuPy uses CUDA-related libraries including cuBLAS, cuDNN, cuRand, cuSolver, cuSPARSE, cuFFT and NCCL to make full use of the GPU architecture. 

 

CuPy's interface is highly compatible with NumPy; in most cases it can be used as a drop-in replacement.

All you need to do is just replace numpywith cupyin your Python code. 

Basics of CuPy (Tutorial)is usefull to learn first step of CuPy. 
CuPy supports various methods, indexing, data types, broadcasting and more. 

Comparison Table (Reference Manual)shows a list of NumPy / SciPy APIs and its corresponding CuPy implementations.

>>> 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)

You can easily make a custom CUDA kernel if you want to make your code run faster, requiring only a small code snippet of C++.

CuPy automatically wraps and compiles it to make a CUDA binary.

Compiled binaries are cached and reused in subsequent runs.

Please read User-Defined Kernels (Tutorial). 
And, you can also use raw CUDA kernel via Raw modules (Tutorial).

>>> 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)

你可能感兴趣的:(CuPy)