Pytorch with nvidia-docker on Ubuntu 18.04


A new computer with 2080 Ti just joint us. This time we try to use nvidia-docker.

We all loved docker continer don’t we? A Docker Tutorial for Beginners


  • Ubuntu 16.04 or later
  • NVIDIA GPU(s) that support CUDA

Tips for LVM

If you install your Ubuntu with LVM, extend the LVM partition before anything else

$ sudo lvm
lvm> lvextend -l +100%FREE /dev/ubuntu-vg/ubuntu-lv
lvm> exit
$ sudo resize2fs /dev/ubuntu-vg/ubuntu-lv

The path /dev/ubuntu-vg/ubuntu-lv is physical device for my root, check your path with sudo fdisk -l.

Install GPU driver

Since CUDA is with the nvidia-docker all we need to do is to install GPU driver.

C++ compiler and other related tools are required to finish the installation, and I am sooooo L-A-Z-Y- that I just install everything I need to build anything.

$ sudo apt-get install build-essential

Find your driver from NVIDIA download center

Screen capture for downloading driver

You will need to answer few question during installation, so don’t leave the screen for too long.

Use docker for the project

Suppose you have installed docker, I choosed to install docker during installation of Ubuntu so it is come with my system… If you need to install it manually, check out official instructions Install Docker Engine on Ubuntu

It will be a good idea to use separated Docker for each project. It is easy to set up and run.

$ sudo docker run --gpus all nvidia/cuda nvidia-smi

I can see it printed the status for all my GPUs.

Use Pytorch docker

$ sudo docker run --gpus all pytorch/pytorch nvidia-smi

It works too, great!

Now mount our data and source code for a trail…

$ sudo docker run --gpus all --rm -v /home/user/code:/workspace -v /home/user/data:/data -v /home/user/outputs:/outputs pytorch/pytorch

Well done!