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