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.
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 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.
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.
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
The output of this command will be the version of PyTorch that is used in the container. Expected result is as show below.