Pytorch+cuda installation script for Ubuntu 20.04 LTS WSL

install.bash

#!/bin/bash

# Installs Pytorch development environment
# on Ubuntu 20.04 LTS (WSL2) host machine
# First make sure that Windows GPU driver and WSL2 installations are ok,
# then launch WSL, create a new user and execute commands.
# Reference: https://docs.nvidia.com/cuda/wsl-user-guide/index.html

# 1) Install cuda stuff
sudo apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub
sudo sh -c 'echo "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64 /" > /etc/apt/sources.list.d/cuda.list'
sudo apt-get update
sudo apt-get install -y cuda-toolkit-11-0

# 2) Verify that GPU can be found (optional step)
cd /usr/local/cuda/samples/1_Utilities/deviceQuery/
sudo make && ./deviceQuery && cd ~/

# 3) Install docker
curl https://get.docker.com | sh
sudo usermod -aG docker $USER

# 4) Install NVIDIA Container Toolkit
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
curl -s -L https://nvidia.github.io/libnvidia-container/experimental/$distribution/libnvidia-container-experimental.list | sudo tee /etc/apt/sources.list.d/libnvidia-container-experimental.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2

# 5) Install Nvidia seasoned Pytorch docker container
sudo service docker start
sudo docker pull nvcr.io/nvidia/pytorch:20.11-py3

# 6) Launch the container
mkdir data
su - $USER # Reload docker group permissions, to be able to run docker
docker run --gpus all -it --rm -v /home/$USER/data:/workspace/data nvcr.io/nvidia/pytorch:20.11-py3

Additional optional steps:

• Install miniconda for managing environments and staying sane with CUDA
• Symlink needed Windows folders (e.g. ln -s /mnt/c/Data/ ~/data)
• Define CUDA related variables in ~/.bashrc (https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html#post-installation-actions)
• Install CUDNN (https://docs.nvidia.com/deeplearning/cudnn/install-guide/index.html#installlinux-deb)

If you plan to run GUI applications on WSL (e.g. matplotlib):
https://techcommunity.microsoft.com/t5/windows-dev-appconsult/running-wsl-gui-apps-on-windows-10/ba-p/1493242

Install jupyter with conda:
$ conda install jupyter
$ conda install -c miniconda ipykernel # this allows to use conda-env with notebook

Create env and then launch notebook with it:
$ python -m ipykernel install --user --name=yolo
Launch without browser: $ jupyter notebook --port=8889 --no-browser
(https://medium.com/@nrk25693/how-to-add-your-conda-environment-to-your-jupyter-notebook-in-just-4-steps-abeab8b8d084)