Machine Learning consulting company should also be adroit at software engineering, rigth? In this post I'll show how to prepare Docker container able to run 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 one-line run command. Sample code is using Keras with TensorFlow backend.

Prerequisites

  1. Already prepared Keras NN. We'll save it to HDF5 file (here you can find more info).
  2. Docker installed (instructions)

Code

For simplicity, we'll be using using well known example - CIFAR10 classification.

Firstly, we have to train our model and save it for later use. Here I'll show just relevant fragment - how to save model to .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 a lot:

# 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, we saved our model. Because 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 saved model and running predictions:

# 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 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 we did here? Let's make a 3 step summary

  1. We trained our model and saved it to model.h5 file
  2. We created Docker Image with our model, keras, tensorflow and all the stuff needed to run our prediction, as well 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 link to repository and 2-line usage instructions.

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