Skip to main content
Join the official 2019 Python Developers SurveyStart the survey!

A small library for managing deep learning models, hyper parameters and datasets

Project description

Zookeeper

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.

Installation

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

Registry

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):
    @property
    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)


@registry.register_preprocess("mnist")
class default(ImageClassification):
    def inputs(self, data):
        return super().inputs(data) / 255

Models

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

@registry.register_model
def cnn(hp, input_shape, num_classes):
    return tf.keras.models.Sequential(
        [
            tf.keras.layers.Conv2D(
                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.Flatten(),
            tf.keras.layers.Dense(hp.filters[3], activation=hp.activation),
            tf.keras.layers.Dense(num_classes, activation="softmax"),
        ]
    )

Hyperparameters

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

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

    @property
    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:

@cli.command()
@build_train()
def train(build_model, dataset, hparams, output_dir):
    """Start model training."""
    model = build_model(hparams, **dataset.preprocessing.kwargs)
    model.compile(
        optimizer=hparams.optimizer,
        loss="categorical_crossentropy",
        metrics=["categorical_accuracy", "top_k_categorical_accuracy"],
    )

    model.fit(
        dataset.train_data(hparams.batch_size),
        steps_per_epoch=dataset.train_examples // hparams.batch_size,
        validation_data=dataset.validation_data(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__":
    cli()

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.

Usage

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

python examples/train.py --help
Usage: train.py [OPTIONS] COMMAND [ARGS]...

Options:
  --help  Show this message and exit.

Commands:
  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.py 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.py train cnn --dataset mnist --hparams-set basic --hparams batch_size=32,actvation=relu

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for zookeeper, version 0.5.5
Filename, size File type Python version Upload date Hashes
Filename, size zookeeper-0.5.5-py3-none-any.whl (19.8 kB) File type Wheel Python version py3 Upload date Hashes View hashes
Filename, size zookeeper-0.5.5.tar.gz (15.3 kB) File type Source Python version None Upload date Hashes View hashes

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN SignalFx SignalFx Supporter DigiCert DigiCert EV certificate StatusPage StatusPage Status page