Docker Deep Learning – GPU-Accelerated Keras

Note: This article was originally published on November 3, 2017 and has been migrated from our previous blog. Some details — tools, libraries, benchmarks, industry context — may be outdated. For our latest perspective, see our recent posts.

Machine Learning consulting companies should also be adept at software engineering, right? In this post, I’ll show you how to prepare a Docker container able to run an already trained Neural Network (NN). It can be helpful if you want to redistribute your work to multiple machines or send it to a client, along with a one-line run command. The sample code is using Keras with TensorFlow backend.

Prerequisites

1. You have already prepared Keras NN. We’ll save it to HDF5 file (here you can find more info).
2. You have installed Docker – (instructions)

Code

For simplicity, we’ll be using using a well known example – CIFAR10 classification.

First, we have to train our model and save it for later use. Here I’ll show just a relevant fragment – how to save the model to an .h5 file, because you probably have your own code that you want to distribute. Note that we probably want to run this in the cloud or on a computer with a good GPU card, so we don’t need to wait too long:

# fragment of learn.py
# whole file can be found in repository
# links below
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
model = Sequential()
model.add(...)

# layers omitted for clarity

model.compile(...)
model.fit(...)

# here we actually save trained model
model.save('model.h5')

Ok, so we have saved our model. Because the model is ready, we don’t need GPU support in our container, which simplifies a lot of things (it’s possible to have GPU support inside docker using nvidia-docker, but it’s more complicated)

A very simple code for loading a saved model and running predictions is:

# predict.py
import argparse
import sys
import os
import glob
import numpy as np

from keras.models import load_model as load_keras_model
from keras.preprocessing.image import img_to_array, load_img

# disable TF debugging info
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

# our saved model file
# may be refactored to be taken from command line
model_filename = 'model.h5'
class_to_name = [
    "Airplane",
    "Automobile",
    "Bird",
    "Cat",
    "Deer",
    "Dog",
    "Frog",
    "Horse",
    "Ship",
    "Truck"
]

def get_filenames():
    parser = argparse.ArgumentParser()
    parser.add_argument('filename', nargs='*', default=['**/*.*'])
    args = parser.parse_args()

    for pattern in args.filename:
        # here we recursively look for input
        # files using provided glob patterns
        for filename in glob.iglob('data/' + pattern, recursive=True):
            yield filename

def load_model():
    if os.path.exists(model_filename):
        return load_keras_model(model_filename)
    else:
        print("File {} not found!".format(model_filename))
        exit()

def load_image(filename):
    img_arr = img_to_array(load_img(filename))
    return np.asarray([img_arr])

def predict(image, model):
    result = np.argmax(model.predict(image))
    return class_to_name[result]

if __name__ == '__main__':
    filenames = get_filenames()
    keras_model = load_model()
    for filename in filenames:
        image = load_image(filename)
        image_class = predict(image, keras_model)
        print("{:30}   {}".format(filename, image_class))

Actual Dockerfile (sorry smartphone users, it’s hard to make that code mobile-friendly). It’s based on the official Dockerfile for Keras:

# Dockerfile
FROM ubuntu:16.04

ENV CONDA_DIR /opt/conda
ENV PATH $CONDA_DIR/bin:$PATH

RUN mkdir -p $CONDA_DIR && \
    echo export PATH=$CONDA_DIR/bin:'$PATH' > /etc/profile.d/conda.sh && \
    apt-get update && \
    apt-get install -y wget git libhdf5-dev g++ graphviz bzip2 && \
    wget --quiet https://repo.continuum.io/miniconda/Miniconda3-4.2.12-Linux-x86_64.sh && \
    echo "c59b3dd3cad550ac7596e0d599b91e75d88826db132e4146030ef471bb434e9a *Miniconda3-4.2.12-Linux-x86_64.sh" | sha256sum -c - && \
    /bin/bash /Miniconda3-4.2.12-Linux-x86_64.sh -f -b -p $CONDA_DIR && \
    rm Miniconda3-4.2.12-Linux-x86_64.sh

ENV NB_USER keras
ENV NB_UID 1000

RUN useradd -m -s /bin/bash -N -u $NB_UID $NB_USER && \
    mkdir -p $CONDA_DIR && \
    chown keras $CONDA_DIR -R && \
    mkdir -p /src && \
    chown keras /src

USER keras

# Python
ARG python_version=3.5
ENV KERAS_BACKEND=tensorflow

RUN conda install -y python=${python_version} && \
    pip install --upgrade pip && \
    pip install tensorflow h5py Pillow && \
    git clone git://github.com/fchollet/keras.git /src && pip install -e /src[tests] && \
    pip install git+git://github.com/fchollet/keras.git && \
    conda clean -yt

ENV PYTHONPATH='/src/:$PYTHONPATH'

WORKDIR /srv
ADD . /srv/

CMD ["python", "-W", "ignore", "predict.py"]

Usage:

# Firstly, let's build our model
# Remember about 'model.h5' file
docker build -t cifar .

# We run created docker image like that
# $PWD should be directory with images
docker run -it --rm -v $PWD:/srv/data cifar python predict.py

# We can also upload image to Docker Registry
# so others can also easily run it
docker tag cifar registry.hub.docker.com/u/valian/cifar
docker push registry.hub.docker.com/u/valian/cifar

All presented code samples can be found in my repository, along with README.md and Makefile to simplify the whole process even more.

🚀 GPU acceleration can take you further. Train a deep neural network to do image recognition tasks or to build a chatbot.

Explanation

What exactly did we do here? Let’s make a 3 step summary:

1. We trained our model and saved it to a model.h5 file
2. We created a Docker Image with our model, keras, tensorflow and all the stuff needed to run our prediction, as well as with file `predict.py` which loads input data (in our case, images) and outputs predictions
3. Now we can distribute “executable box” with our Neural Network model either by Docker Registry or by giving a link to the repository and 2-line usage instructions.