Deep learning with containers. Part 1 | by Alexander Visheratin | Towards Data Science

Containers are lightweight software packages that run in isolation on the host computing environment. Using containerized environments for the research is helpful because they are easy to use, reproducible, and keep your primary system clean. You can easily set up a fully functioning multi-container application using Docker Compose. All scripts required for this are available here.
Update 02/2021: the second part of the series is available here.
When doing intensive research, sometimes it becomes hard to maintain what is going on around you. Tons of files with data, dozens of libraries that you don’t use but not sure enough to wipe them out, temporal files saved somewhere in your home directory, you name it. Virtual environments help but not much. You still have to manage them and might end up in the same chaotic position but multiplied by the environments’ number. While setting up my laptop for working on deep NLP models, I came up with the way of creating easy to use and fully disposable environments for experimenting — containers. They allow me to set up ready-to-go Jupyter notebooks with Tensorboard tracking with only one command. And with another command I can pull this all down and not worry about something that is left in the memory or on the hard drive. In this post I will describe how to set up and deploy containers for deep learning in your system, and in the next one I will show how to extend the basic images for a specific task — siamese learning with Hugging Face Transformers.
The stack that I use is the following:
Docker. Platform and software for creating containers. De-facto standard tool in the field.
PyTorch with GPU support. My framework of choice, it’s just awesome.
Jupyter Notebook. Software that enables interactive development in the browser.
Tensorboard. Good starting point for experiments tracking. From that you can switch to Allegro AI or Weights & Biases.
Hugging Face. Simply the best way you can get state-of-the-art models based on Transformers and other architectures.
If you are not familiar with the concept of containers, it would be a good idea to read some introductory articles (e.g. from Docker or Google) and go through tutorials that cover the basics.
We will now create two containers — one for Jupyter Notebook and another for Tensorboard. We start with the Tensorboard container because it is very small and easy to understand. Here is the full Dockerfile for it:
# Use small image with Python pre-installed.FROM python:3.8-slim-buster# Install tensorboardRUN pip3 install tensorboard
For running and operating Tensorboard we need nothing except its package. You can now save this file as tensorboard.Dockerfile, build the container (docker build -t dl-tensorboard:latest -f tensorboard.Dockerfile .) and run it interactively with the command docker run --rm -it dl-tensorboard:latest /bin/bash, the command tensorboard will be available to run. But it will not run on its own, it needs a directory where the data about experimental runs is stored. We shall deal with it later on the services setup stage.
The next part is building an image with the Jupyter Notebook. There are quite a lot nuances here. First of all, we are going to work with deep learning, so we need a GPU support in our containers. Luckily, NVIDIA got us covered with their Container Toolkit. After several simple steps you get GPU-enabled Docker containers at your disposal. Another very useful project from NVIDIA is a set of pre-built images with CUDA installed. A couple versions of Ubuntu ago, it was a huge pain to install NVIDIA drivers and CUDA. Now it is much easier, but with these images you don’t need to do a thing related to CUDA setup, which is pretty cool.
By the time of writing this post, PyTorch supports CUDA version 10.2, so we will use NVIDIA base image with CUDA 10.2. Additionally, NVIDIA provides three flavors of their images — base, runtime, and devel. Our option is runtime because it includes cuDNN.
The next interesting point is that we are going to perform a multi-stage build of our image. It means that during the build, we will have a base image that encapsulates all the logic related to setting things up, and the final image that adds finishing touches depending on the task in hand, e.g. install additional packages or copy required files. This logic will allow us to have a base image resting in the Docker cache so there will be no need to, for example, download PyTorch every time when we add a new package. Without a doubt, we can achieve the same without a multi-stage build by just adding new lines in Dockerfile after the existing and not messing up the order of already added lines (because of the way how Docker cache works), but this approach gives a division of responsibilities, ease of operation, and lays the basis for more advances scenarios that we will cover in next posts.
I would like to clarify a couple more points before we jump into the Dockerfile for a Jupyter image. First, since we use NVIDIA CUDA image as a base, we will need to install Python and some required tools. Second, due to the known issue it is not possible to run a Notebook binary in the container as-is. Some time ago it would require setting up Tini in the Dockerfile, as described in the link. But now Tini is built into the Docker so there is no need for any additional setup, just add --init flag when running the container. A bit later we shall see how to incorporate this to the Docker compose.
Finally, we can have a look at the actual Dockerfile for the Jupyter Notebook image:
# NVIDIA CUDA image as a base# We also mark this image as "jupyter-base" so we could use it by nameFROM nvidia/cuda:10.2-runtime AS jupyter-baseWORKDIR /# Install Python and its toolsRUN apt update && apt install -y --no-install-recommends \    git \    build-essential \    python3-dev \    python3-pip \    python3-setuptoolsRUN pip3 -q install pip --upgrade# Install all basic packagesRUN pip3 install \    # Jupyter itself    jupyter \    # Numpy and Pandas are required a-priori    numpy pandas \    # PyTorch with CUDA 10.2 support and Torchvision    torch torchvision \    # Upgraded version of Tensorboard with more features    tensorboardX# Here we use a base image by its name - "jupyter-base"FROM jupyter-base# Install additional packagesRUN pip3 install \    # Hugging Face Transformers    transformers \    # Progress bar to track experiments    barbar
Having in mind everything that was described before, the script is quite straightforward — take NVIDIA CUDA image, install Python, install basic packages, use this all as a basis for the final image with some additional packages that depend on the specific problem.
The final step of creating the environment is combining both images into the complete multi-container application with the help of Docker Compose. This tool allows to create a YAML file that contains a description of services that form the application. This is a nice tool that is mainly focused on a single-host deployments, which is exactly our use case. One majoe drawback of Compose at the moment of writing this post is the lack of GPU support, despite the fact that Docker and its API have implemented this feature. But there is a workaround that enables usage of GPUs in Compose services. I have integrated it into the latest release of Compose and plan to maintain it until GPU and other device requests are supported. With this regard, you need to install this fork of Compose to be able to run the services.
Let’s look into the docker-compose.yml file:
version: "3.8"services:    tensorboard:        image: dl-tensorboard        build:            context: ./            dockerfile: tensorboard.Dockerfile        ports:            - ${TENSORBOARD_PORT}:${TENSORBOARD_PORT}        volumes:            - ${ROOT_DIR}:/jupyter        command:            [                "tensorboard",                "--logdir=${TENSORBOARD_DIR}",                "--port=${TENSORBOARD_PORT}",                "--bind_all",            ]    jupyter-server:        image: dl-jupyter        init: true        build:            context: ./            dockerfile: jupyter.Dockerfile        device_requests:            - capabilities:                - "gpu"        env_file: ./.env        ports:            - ${JUPYTER_PORT}:${JUPYTER_PORT}        volumes:            - ${ROOT_DIR}:/jupyter        command:            [                "jupyter",                "notebook",                "--ip=*",                "--port=${JUPYTER_PORT}",                "--allow-root",                "--notebook-dir=${JUPYTER_DIR}",                '--NotebookApp.token=${JUPYTER_TOKEN}'            ]
There is a lot to unwrap here. First of all, there are two services defined — tensorboard and jupyter-server. The image section defines the name of the image that will be built for the service. The build section defines a path to the directory or a repository with the Dockerfile that will be used for building the image. In the ports section we specify what ports from the host machine will be exposed to the container. And here the things become interesting – instead of numeric values we see names of environment variables. Docker compose takes these variables from a special file named .env where the user lists all variables that have to be used during compose. This allows to have a single place where you store all parts of the compose script that can differ from environment to environment. For example, the port that is used for the Tensorboard on one machine can be occupied on the other, and this issue can be easily fixed by changing the value of TENSORBOARD_PORT variable in the .env file.
For our scenario, .env file looks like this:
ROOT_DIR=/path/to/application/root/directory # path on the hostJUPYTER_PORT=42065JUPYTER_TOKEN=egmd5hashofsomepasswordJUPYTER_DIR=/jupyter/notebooks               # path in the containerTENSORBOARD_PORT=45175TENSORBOARD_DIR=/jupyter/runs                # path in the container
ROOT_DIR is a directory that will store all information that is needed by our application (notebooks, data, runs history). As you can see from the volumes section in docker-compose.yml, this directory is mounted to both containers as /jupyter, which explains why both JUPYTER_DIR and TENSORBOARD_DIR variables have this prefix. Other variables in the file define the ports to be exposed and a token for Jupyter Notebook that is used for authentication.
The command section of the service definition contains a set of arguments that combined form a command that will be executed on the container start-up. In order to fix the described above issue with a Notebook binary, we add init: true for the Jupyter service. We also pass environment variables from .env file to the container with the Jupyter Notebook using the section env_file. This is because we will make use of these variables in the notebooks, details on this will be in the next post. And last but not the least feature of the services definition is a device_requests section. This is exactly the workaround that enables a GPU to be available inside the container. Without it you won’t be able to enjoy the power of your top-notch NVIDIA card in the multi-container application.
In the end, we get the following structure of the working directory:
├── .env                     # File with environment variables.├── docker-compose.yml       # Compose file with services.├── jupyter.Dockerfile       # Dockerfile for the Jupyter image.├── tensorboard.Dockerfile   # Dockerfile for the Tensorboard image.
With all set and done, you can finally run docker-compose up -d inside our working directory. After downloading all base images and installing all required packages both containers should be up and running. The first start takes quite a long time but all next will be very fast – even if you install new packages in the final image, all heavylifting with downloading and installing base packages will be skipped. And when you are done with your work, just run docker-compose down and all containers will disappear and leave your system clean and ready for other fascinating stuff, like developing distributed systems in Go or running complex physics simulations in COMSOL.
In this post, we looked into an example of how to set up a containerized environment with Jupyter Notebook, PyTorch and Tensorboard that can be used for the deep learning research. In the next post, we will dive into how this environment can actually be used for a real use case — siamese learning for a few-shot classification of the news. Key outtakes from this post:
Use multi-stage builds, where the base image does one-time heavy things, and the final image tailors the environment to your needs. You know, like fine-tuning.Docker Compose does not support binding GPU to containers. But if you really want, it does.
Use multi-stage builds, where the base image does one-time heavy things, and the final image tailors the environment to your needs. You know, like fine-tuning.
Docker Compose does not support binding GPU to containers. But if you really want, it does.