Running PyTorch in Docker Container

Posted On May 16, 2020

Hi, does it that frustrate you that it is very complicated to setup your environment when you want to try out a new framework. Luckily, there is a solution for that by using docker container. In this post we will try out how to train neural network using PyTorch network inside docker container.

Environment

I am going to assume that you use the following system configuration:

• Operating System : Ubuntu 20.04
• CPU : Intel i7
• GPU : Nvidia GTX 1080 Ti

Installing Docker

Installing docker is as simple as opening up a terminal and run the following command:

sudo apt install docker.io

Once installation is successful, we need to add current user to the docker user group to allow current user to run docker without needing to call sudo. This is done by creating a docker user group.

sudo groupadd docker
sudo usermod -aG docker $USER
newgrp docker

Installing Nvidia-docker

Training neural network requires huge computational resource, this can be accelerated using GPU. We can enable support for Nvidia GPUs by installing Nvidia-docker. The full instruction on how to do this can be found here. [https://github.com/NVIDIA/nvidia-docker]

Below are the required commands as described in the above link.

# Add the package repositories
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

sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

Running PyTorch container

Once docker is setup properly, we can run the container using the following commands:

docker run --rm --name pytorch --gpus all -it pytorch/pytorch:1.5-cuda10.1-cudnn7-devel

The above command will run a new container based on the PyTorch image specified by “pytorch/pytorch:1.5-cuda10.1-cudnn7-devel”. This will run container with PyTorch 1.5 with cuda 10.1 setup in the container. The complete list of images can be found here : https://hub.docker.com/r/pytorch/pytorch.

We will now test if PyTorch is correctly setup in the container. From within the container, try to run python and run the following commands

python
import torch
print(torch.__version__)

The output of this command will be the version of PyTorch that is used in the container. Expected result is as show below.