Skip to content
A small library for managing deep learning models, hyper parameters and datasets
Branch: master
Clone or download
Fetching latest commit…
Cannot retrieve the latest commit at this time.
Type Name Latest commit message Commit time
Failed to load latest commit information.
examples Update example and readme Aug 8, 2019
zookeeper Remove `plot()` command test. Aug 16, 2019
.gitignore Replace tf.contrib.HParams with custom class May 29, 2019


Azure DevOps builds Azure DevOps coverage PyPI - Python Version PyPI PyPI - License Code style: black Join the community on Spectrum

A small library for managing deep learning models, hyper parameters and datasets designed to make training deep learning models easy and reproducible.

Getting Started

Zookeeper allows you to build command line interfaces for training deep learning models with very little boiler plate using click and TensorFlow Datasets. It helps you structure your machine learning projects in a framework agnostic and effective way. Zookeeper is heavily inspired by Tensor2Tensor and Fairseq but is designed to be used as a library making it lightweight and very flexible.


pip install zookeeper
pip install colorama  # optional for colored console output


Zookeeper keeps track of data preprocessing, models and hyperparameters to allow you to reference them by name from the commandline.

Datasets and Preprocessing

TensorFlow Datasets provides many popular datasets that can be downloaded automatically. In the following we will use MNIST and define a default preprocessing for the images that scales the image to [0, 1] and uses one-hot encoding for the class labels:

import tensorflow as tf
from zookeeper import cli, build_train, HParams, registry, Preprocessing

class ImageClassification(Preprocessing):
    def kwargs(self):
        return {
            "input_shape": self.features["image"].shape,
            "num_classes": self.features["label"].num_classes,

    def inputs(self, data):
        return tf.cast(data["image"], tf.float32)

    def outputs(self, data):
        return tf.one_hot(data["label"], self.features["label"].num_classes)

class default(ImageClassification):
    def inputs(self, data):
        return super().inputs(data) / 255


Next we will register a model called cnn. We will use the Keras API for this:

def cnn(hp, input_shape, num_classes):
    return tf.keras.models.Sequential(
                hp.filters[0], (3, 3), activation=hp.activation, input_shape=input_shape
            tf.keras.layers.MaxPooling2D((2, 2)),
            tf.keras.layers.Conv2D(hp.filters[1], (3, 3), activation=hp.activation),
            tf.keras.layers.MaxPooling2D((2, 2)),
            tf.keras.layers.Conv2D(hp.filters[2], (3, 3), activation=hp.activation),
            tf.keras.layers.Dense(hp.filters[3], activation=hp.activation),
            tf.keras.layers.Dense(num_classes, activation="softmax"),


For each model we can register one or more hyperparameters sets that will be passed to the model function when called:

class basic(HParams):
    activation = "relu"
    batch_size = 32
    filters = [64, 64, 64, 64]
    learning_rate = 1e-3

    def optimizer(self):
        return tf.keras.optimizers.Adam(self.learning_rate)

Training loop

To train the models registered above we will need to write a custom training loop. Zookeeper will then tie everything together:

def train(build_model, dataset, hparams, output_dir):
    """Start model training."""
    model = build_model(hparams, **dataset.preprocessing.kwargs)
        metrics=["categorical_accuracy", "top_k_categorical_accuracy"],
        steps_per_epoch=dataset.train_examples // hparams.batch_size,
        validation_steps=dataset.validation_examples // hparams.batch_size,

This will register Click command called train which can be executed from the command line.

Command Line Interface

To make the file we just created executable we will add the following lines at the bottom:

if __name__ == "__main__":

If you want to register your models in separate files, make sure to import them before calling cli to allow zookeeper to properly register them. To install your CLI as a executable command checkout the setuptools integration of Click.


Zookeeper already ships with prepare, plot, and tensorboard commands, but now also includes the train command we created above:

python examples/ --help

  --help  Show this message and exit.

  install-completion  Install shell completion.
  plot                Plot data examples.
  prepare             Downloads and prepares datasets for reading.
  tensorboard         Start TensorBoard to monitor model training.
  train               Start model training.

To train the model we just registered run:

python examples/ train cnn --dataset mnist --hparams-set basic --hparams batch_size=64

Multiple arguments are seperated by a comma, and strings should be passed without quotion marks:

python examples/ train cnn --dataset mnist --hparams-set basic --hparams batch_size=32,actvation=relu
You can’t perform that action at this time.