Packaging your PyTorch project in Docker
Have you ever struggled setting up your deep learning project on a new machine? In this article, we will discover how to package your project inside Docker. You won’t even need to install CUDA and cuDNN! Thanks to NVIDIA-docker, you can harness the power of your machine’s GPU at no performance cost. Training and inference will be as fast as on your machine ✨🚤
With 2 simple commands, you’ll be able to run your project inside a Docker container. It works on any machine, as long as Docker and Nvidia-Docker are installed. No more time spent meddling with requirements or obscure dependencies.
Here is a small summary of how the project will look like:
(This article assumes the use of GPU, but it is entirely optional)
Follow the instructions here to install NVIDIA-docker (it takes a few minutes). Then, create a file named
Dockerfile at the root of your project. In it, copy and paste the following:
FROM pytorch/pytorch:1.9.0-cuda10.2-cudnn7-runtime COPY ./requirements.txt /install/requirements.txt RUN pip install -r /install/requirements.txt WORKDIR /code
Only two commands to know:
Go in your project directory, and type:
sudo docker build -t my-project .
(Note the dot at the end). It will download all the dependencies specified in the
Dockerfile you just created.
When it is ready, run the following command:
sudo docker run -it --gpus "device=1" -v $(pwd):/code my-project:latest python3 train.py
You may run into an error. It can be that a version is not set correctly or that a dependency is missing. For instance, attempting to
import cv2 will yield the following error:
ImportError: libGL.so.1: cannot open shared object file: No such file or directory
Thankfully, this issue is easy to resolve. Update the command above to the following one:
sudo docker run -it --gpus "device=1" -v $(pwd):/code my-project:latest bash
This opens a shell inside the docker container. Now that we are running inside it run
Use your extensive Google-fu to find out the missing dependency and install it. Note down what you are typing, as we will migrate those precious commands to the
Dockerfile to add the missing dependencies. Prefix any command with RUN. For instance, if you’d like to add support for OpenCV, the best way is to let
apt install all dependencies by adding the following line in 2nd:
RUN apt-get update && apt-get install -y python3-opencv
Now, we are ready for one last
docker build . and we are good to go! Type the
run command from above, and your training should start right away 🚀
Tips and other useful commands
- The command
docker build .will attempt to copy all files in your directory to create a “Context”. This might be slow if you have a lot of data in this folder. Therefore we recommend creating a
.dockerignorefile with your models and other large files.
- Docker can leave stale containers hanging on your disk. Run
sudo docker system prune -aif too much disk space is used.
- Each line in the Dockerfile has its own cache. It means that running
docker builda second time is usually much faster if you did not change many lines
- One benefit is that you don’t need to install CUDA or cuDNN on the host system. NVIDIA drivers for your GPU still need to be installed