Introduction
Estimator
A predictable way to train your deep learning model
Training deep learning model in tensorflow should be easy and fast. This library introduces a predictable object-oriented approach to achieve booth simplicity and granularity while preserving customization.
The aim of this library is to allows the developer to just define a model and an input pipeline in order to train, evaluate and test their model without writing any train loop and reinvent the wheel every time.
Installation
The package is available on pip:
pip install TODO
or you can clone this directory and manually place the estimator
folder in your project.
Then it can be imported in your script
from estimator.Estimator import Estimator
Motivation
TensorFlow provides an Estimator implementation that is more a black box than a library. It cannot be easily extended and customize for this reasons I decided to implement my own.
Example
It follows a very simple example in which we create a basic feed forward neural network and we train it with random numbers
from estimator.Estimator import Estimator, Mode
import numpy as np
import tensorflow as tf
EPOCHS = 100
# create the dataset
train_data = (np.random.rand(1000,2),np.random.rand(1000,1))
test_data = (np.random.rand(100,2),np.random.rand(100,1))
# define the input function
input_fn = Estimator.create_input_fn(input_shape=[None,2],output_shape=[None,1])
# define the model builder
def model_builder(x, y, config):
# config is a dictionary that can be passed to the estimator
net = tf.layers.dense(x, 16, activation=tf.nn.relu)
predictions = tf.layers.dense(net, 1, activation=tf.nn.sigmoid)
loss = tf.losses.mean_squared_error(labels=y, predictions=predictions)
train_step = tf.train.AdamOptimizer(0.01).minimize(loss)
# it must return a dictionary contain the operation to train, predict and evaluate
return {
Mode.TRAIN: {'train_step': train_step },
Mode.PREDICT: { 'predictions': predictions },
Mode.EVAL: { 'loss': loss } # used as metrics
}
estimator = Estimator(model_builder, input_fn)
# we can define a batch size before train, default is one
estimator.train(data=train_data, epochs=EPOCHS, batch_size=64)
res = estimator.evaluate(data=test_data)
print(res)
pred = estimator.predict(np.array([[2,1]]))
print(pred)
Data creation
In this example we fist create two random dataset using numpy
train_data = (np.random.rand(1000,2),np.random.rand(1000,1))
test_data = (np.random.rand(100,2),np.random.rand(100,1))
Input pipeline
Then we define an input_fn
by calling an utily function from the Estimator
class that takes the input and output shape.
input_fn = Estimator.create_input_fn(input_shape=[None,2],output_shape=[None,1])
You must manually specify the shapes since, internally, the Estimator will create a generic iterator that will switch between the train Dataset and the eval Dataset.
You must pass None
as the first dimension since the first dimension is used later for batching.
Model creation
def model_builder(x, y, config):
# config is a dictionary that can be passed to the estimator
net = tf.layers.dense(x, 16, activation=tf.nn.relu)
predictions = tf.layers.dense(net, 1, activation=tf.nn.sigmoid)
loss = tf.losses.mean_squared_error(labels=y, predictions=predictions)
train_step = tf.train.AdamOptimizer(0.01).minimize(loss)
# it must return a dictionary contain the operation to train, predict and evaluate
return {
Mode.TRAIN: {'train_step': train_step },
Mode.PREDICT: { 'predictions': predictions },
Mode.EVAL: { 'loss': loss } # used as metrics
}
The Estimator needs a function that returns a dictionary with the operations to run for each mode. The modes are three:
- Mode.Train: defines the operations that will be run when we call the
.train
method - Mode.Predict: defines the operations that will be run when we call the
.predict
method - Mode.Eval: defines the operations that will be run when we call the
.eval
method. These are your metrics. For example, accuracy
The model is built only once.
Estimator creation
Create an estimator is super easy
estimator = Estimator(model_builder, input_fn)
Train
To train, just call the .train
method with the train data and a batch size as input:
estimator.train(data=train_data, epochs=EPOCHS, batch_size=64)
Evaluate
To evaluate, call the .evaluate
method. You can also pass a batch size, the default is the size of the data input
res = estimator.evaluate(data=test_data)
It will return a dictionary that contains the results of all the operation defined in the field Mode.EVAL
of the dictionary that model_builder
returns. In our case, the output is:
{'loss': [0.08878718]}
Train and Evaluate
You can train and evaluate by calling
estimator.train_and_evaluate(data=train_data, epoches=10, validation=val_data, batch_size=64, batch_size_eval=32, every=2)
Predict
To predict, call the .predict
method with an input.
pred = estimator.predict(np.array([[2,1]]))
print(pred)
It will return a dictionary that contains the results of all the operation defined in the field Mode.PREDICT
of the dictionary that model_builder
returns. In our case, the output is:
{'predictions': array([[0.61255413]], dtype=float32)}
Customisation
Hooks can be attach the Estimator in order to decorate it with new functionalities. Each hook must subclass the Hook
class.
class Hook:
def before_build_model(self, estimator):
pass
def after_build_model(self, estimator):
pass
def before_run_epoch(self, estimator, epoch, data, batch_n, tot_res):
pass
def after_run_epoch(self, estimator, epoch, data, batch_n, tot_res):
pass
def before_run_batch(self, estimator, res, i):
pass
def after_run_batch(self, estimator, res, i, tot_res):
pass
This class exposes a wide range of function that will be called by the estimator. For example, a very simple logger can be defined as follow:
class Logger(Hook):
def after_run_epoch(self, estimator, epoch, data, batch_n, tot_res):
# tot_res is an array with the result for each batch, in this case { 'loss' : [...] }
# take the mean of each metric key
mean_res = {k: np.mean(tot_res[k]) for k in estimator.metrics[Mode.EVAL].keys()}
print(mean_res)
Then it can be passed to the Estimator as argument
estimator = Estimator(model_builder, input_fn, hooks=[Logger()])
# we can define a batch size before train, default is one
estimator.train(data=train_data, epochs=EPOCHS, batch_size=64)
res = estimator.evaluate(data=test_data)
pred = estimator.predict(np.array([[2,1]]))
print(pred)
The output is:
...
{'loss': 0.08635884}
{'loss': 0.084242225}
{'loss': 0.08413396}
{'loss': 0.082261935}
{'loss': 0.08626987}
{'loss': 0.08318842}
{'loss': 0.08242814}
{'loss': 0.08449897}
{'loss': 0.08369716}
{'loss': 0.0927686}
{'predictions': array([[0.45650724]], dtype=float32)}
A full example can be found in example/hook/main.py