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!


OpenCV Workflow in Jupyter Notebook

OpenCV is great for classical computer vision, and Jupyter Notebooks are great for prototyping. It would make sense to use them both, but the “cv2.imshow” cannot work inside a Jupyter Notebook, since it requires handling an event loop. Okay, you can always convert back the image to a RGB Numpy array and show it using Matplotlib or maybe Pillow. I prefer to avoid the additional lines of code and mental context switch using a small neat package called “cv2-plt-imshow”.

pip install cv2-plt-imshow

Then, reading and showing image in the Notebook is as it should be!

import cv2
from cv2_plt_imshow import cv2_plt_imshow
img = cv2.imread("10000.jpg")

The documentation is here: