Docker Deep Learning - GPU-accelerated Keras

November 3, 2017

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.


  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)


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
# whole file can be found in repository
# links below
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
model = Sequential()

# layers omitted for clarity


# here we actually save trained model'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:

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 = [

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)
        print("File {} not found!".format(model_filename))

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

RUN mkdir -p $CONDA_DIR && \
    echo export PATH=$CONDA_DIR/bin:'$PATH' > /etc/profile.d/ && \
    apt-get update && \
    apt-get install -y wget git libhdf5-dev g++ graphviz bzip2 && \
    wget --quiet && \
    echo "c59b3dd3cad550ac7596e0d599b91e75d88826db132e4146030ef471bb434e9a *" | sha256sum -c - && \
    /bin/bash / -f -b -p $CONDA_DIR && \


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

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


ADD . /srv/

CMD ["python", "-W", "ignore", ""]


# 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

# We can also upload image to Docker Registry
# so others can also easily run it
docker tag cifar
docker push

All presented code samples can be found in my repository, along with 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.


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 `` 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 link to repository and 2-line usage instructions.

Eager to read more articles like this one? Follow us on Twitter!

[email protected]
linkedin facebook pinterest youtube rss twitter instagram facebook-blank rss-blank linkedin-blank pinterest youtube twitter instagram