Following instructions assume you have a CUDA compatible GPU with at least 8GB VRAM (GTX1070 or better) as part of hardware.
The installation of PyTorch GPU in Ubuntu 20.04 can be summarized in the following points,
• Install CUDA by installing nvidia-cuda-toolkit.
• Install the cuDNN version compatible with CUDA.
• Export CUDA environment variables.
First open a terminal and run
$ sudo apt install nvidia-cuda-toolkit
which directly installs the latest version of CUDA in Ubuntu. After installing CUDA, run
$ nvcc -V
You will get an output similar to the following to verify if you had a successful installation,
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Sun_Jul_28_19:07:16_PDT_2019
Cuda compilation tools, release 10.1, V10.1.243
Note the CUDA version above release 10.1
After above step, visit - https://developer.nvidia.com/rdp/cudnn-download - and download the CUDNN package that matches your CUDA version which is highlighted above. Once downloaded run
$ tar -xvzf cudnn-10.1-linux-x64-v7.6.5.32.tgz
Note that the package name might vary in your case. Now move the extracted packages like so
$ sudo cp cuda/include/cudnn.h /usr/lib/cuda/include/
$ sudo cp cuda/lib64/libcudnn* /usr/lib/cuda/lib64/
Set the file permissions of cuDNN,
$ sudo chmod a+r /usr/lib/cuda/include/cudnn.h
$ sudo chmod a+r /usr/lib/cuda/lib64/libcudnn*
The CUDA environment variables are needed by PyTorch for GPU support. To set them, we need to append them to ~/.bashrc
file by running,
$ echo 'export LD_LIBRARY_PATH=/usr/lib/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
$ echo 'export LD_LIBRARY_PATH=/usr/lib/cuda/include:$LD_LIBRARY_PATH' >> ~/.bashrc
Load the exported environment variables by running,
$ source ~/.bashrc
Finally we can check if everything went fine by running
$ nvidia-smi
The installation of PyTorch GPU in Windows 10 can be summarized in the following points,
• Install CUDA
• Install the cuDNN version compatible with CUDA.
• Export CUDA environment variables.
Visit https://developer.nvidia.com/cuda-downloads and download the cuda-toolkit from here
After above step, visit - https://developer.nvidia.com/rdp/cudnn-download - and download the CUDNN package that matches your CUDA version. Once downloaded unzip the cuDNN package.
cudnn-x.x-windows-x64-vx.x.x.x.zip
Copy the following files into the CUDA Toolkit directory.
- Copy
<installpath>\cuda\bin\cudnn*.dll
to
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vx.x\bin
- Copy
<installpath>\cuda\include\cudnn*.h
to
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vx.x\include.
- Copy
<installpath>\cuda\lib\x64\cudnn*.lib
to
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vx.x\lib\x64.
Variable Name: CUDA_PATH
Variable Value: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vx.x=
Finally we can check if everything went fine by running from a new command prompt.
$ nvidia-smi