NNabla CUDA extension package installation using PIP

Note: please refer to the OS specific workflows for the OS specific dependencies setup.

By installing the NNabla CUDA extension package nnabla-ext-cuda, you can accelerate the computation by NVIDIA CUDA GPU (CUDA must be setup on your environment accordingly).

Several pip packages of NNabla CUDA extension are provided for each CUDA version and its corresponding cuDNN version as following.

CUDA vs cuDNN Compatibility

Package name

CUDA version

cuDNN version



7.6(Linux & Win)



8.0(Linux & Win)



8.0(Linux & Win)

The latest CUDA version is always preferred if your GPU accepts.

Currently, for each NNabla CUDA extension package, it may be not compatible with some specific GPUs.

After nnabla-ext-cuda package is installed, you can manually check whether your GPU is usable. For example, you can check GPU with device_id 0 by:

import nnabla_ext.cudnn
device_id = '0'

Above code will run successfully if your GPU is usable, otherwise, an error will be reported.

nnabla-ext-cuda package will also try to check the compatibility of your GPUs automatically when you use ‘cuda’ or ‘cudnn’ extension. By default, it will list and check all gpus in your machine. Error will be reported if there is incompatible card.

You can set environment variable ‘AVAILABLE_GPU_NAMES’ to tell it which GPU is usable, ‘AVAILABLE_GPU_NAMES’ is a white list, GPU in ‘AVAILABLE_GPU_NAMES’ will not cause error. For example, if you think GeForce RTX 3070 and GeForce RTX 3090 are usable, you can set environment variable as following:

export AVAILABLE_GPU_NAMES="GeForce RTX 3070,GeForce RTX 3090"


The following is an example of installing the extension for CUDA 10.2

pip install nnabla-ext-cuda102

and check if all works.

python -c "import nnabla_ext.cuda, nnabla_ext.cudnn"
2018-06-26 15:20:36,085 [nnabla][INFO]: Initializing CPU extension...
2018-06-26 15:20:36,257 [nnabla][INFO]: Initializing CUDA extension...
2018-06-26 15:20:36,257 [nnabla][INFO]: Initializing cuDNN extension...

Note: If you want to make sure the latest version will be installed, try to uninstall previously installed one with pip uninstall -y nnabla nnabla-ext-cuda100 beforehand.

Installation with Multi-GPU supported

Multi-GPU wheel package is only available on python3.6+.

CUDA vs cuDNN Compatibility

You can install as the following.

pip install nnabla
pip install nnabla-ext-cuda100-nccl2-mpi2-1-1

If you already installed NNabla, uninstall all of it, or start from a clean environment which you create using Anaconda, venv.

You should also install OpenMPI and NCCL in addition to CUDA and CuDNN.

If you are using Ubuntu18.04 and choose mpi2.1.1, you can install mpi with following command.

sudo apt install -y --no-install-recommends openmpi-bin libopenmpi-dev

Otherwise, you must install openmpi with following command.(MPIVER=3.1.6 or 2.1.1)

curl -O https://download.open-mpi.org/release/open-mpi/v${MPIVER%.*}/openmpi-${MPIVER}.tar.bz2
tar xvf openmpi-${MPIVER}.tar.bz2
cd openmpi-${MPIVER}
./configure --with-sge
sudo make install


Q. How do I install CUDA?

NNabla CUDA extension requires both CUDA toolkit and cuDNN library. You should select a proper CUDA version according to your CUDA device capability. See the official installation guide. NNabla supports CUDA versions later than 8.0. See the table for the cuDNN compatibility with the specific CUDA versions.

Q. How do I install NCCL

Please visit NCCL, then follow the instruction.

Q. How do I check proper version of cuDNN

Enter the following command:

python -c "import nnabla_ext.cuda, nnabla_ext.cudnn"

If there is a version mismatch on your machine, you can see proper versions in the error message. Following is a sample error message.

[nnabla][INFO]: Initializing CPU extension...
Please install CUDA version 10.2.
  and cuDNN version 8.0
  Or install correct nnabla-ext-cuda for installed version of CUDA/cuDNN.