Getting Ubuntu 20.04 LTS Ready for Pytorch 1.6 with CUDA

Each time I have to re-install Linux I have to go through many blog articles to get up and running. This concise guide is just a quick reminder of what need to be done to get up and running fast.

Installing the latest NVidia drivers.

You can install the drivers by enabling the proprietary drivers for Ubuntu 20.04, and install the latest drivers (as of now, 455).

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt install nvidia-driver-455

Installing CUDA

Installing CUDA will be automatically handled by Anaconda when installing PyTorch. This section can be skipped.

If you need to install CUDA, for instance, in ordert to build Darknet, the process is very straightforward. Just follow the instructions on their website. If you want to cut corners, the instructions are copied here, but I guess this is subject to change.

sudo mv /etc/apt/preferences.d/cuda-repository-pin-600
sudo dpkg -i cuda-repo-ubuntu2004-11-1-local_11.1.1-455.32.00-1_amd64.deb
sudo apt-key add /var/cuda-repo-ubuntu2004-11-1-local/
sudo apt-get update
sudo apt-get -y install cuda

Installing PyTorch

My preference is PyTorch. Installation is straightforward with Anaconda and should be done as recommended on their official site. The only somewhat tricky part is to make sure to choose the PyTorch version compatible with the CUDA version.

conda install pytorch torchvision torchaudio cudatoolkit=11.0 -c pytorch

To make sure everything is working properly, you can start an interactive python shell, import the PyTorch library and check if the GPU is seen.

>>> import torch
>>> torch.cuda.get_device_name(torch.cuda.current_device())
‘GeForce GTX 1070 Ti’

There we go, ready to train neural networks!

Leave a Reply

Your email address will not be published. Required fields are marked *