An Advanced Example of the Tensorflow Estimator Class | by Tijmen Verhulsdonck | Towards Data Science

Estimators were introduced in version 1.3 of the Tensorflow API, and are used to abstract and simplify training, evaluation and prediction. If you haven’t worked with Estimators before I suggest to start by reading this article and get some familiarity as I won’t be covering all of the basics when using estimators. Instead I hope to demystify and clarify some aspects more detailed aspects of using Estimators and switching to Estimators from an existing code-base.
Anyone who has been working with Tensorflow for a long time will know that it used to be a lot more time consuming to setup and use Tensorflow as it is today. Today there are numerous libraries that simplify development such as slim, tflearn and more. These often reduced the code needed to define a network from multiple files and classes to a single function. They also simplified managing the training and evaluation, and to a small extend data preparation and loading. The Estimator class of Tensorflow does not change anything about the network definition but it simplifies and abstracts managing training, evaluation and prediction. It stands out from the other libraries due to it’s low level optimizations, useful abstractions and support from the core Tensorflow dev team.
That is the long story, in short Estimators are faster to run and to implement, simpler (once you get used to them) and well supported.
This article will break down some of the features of the Estimators, using examples from a GitHub project of mine SqueezeNext-Tensorflow. The network implemented in the project is from a paper released in 2018 called “SqueezeNext”. This network is very lightweight and fast due to a novel approach to separable convolutions. The researchers released a caffe version but to facilitate experimentation using the available Tensorflow libraries I recreated the algorithm from the paper in Tensorflow.
The article is setup using the following structure:
Setting Up The Estimator
Data Loading with Estimators and Datasets
Defining Prediction,Training and Evaluation modes
Session hooks and scaffolds
Prediction
In order to setup the Estimator for a certain training and evaluation roster it’s good to first understand how the Estimator setup works. To construct an instance of the Estimator class you can use the following call:
classifier = tf.estimator.Estimator(model_dir=model_dir,                                    model_fn=model_fn,                                    params=params)
In this call “model_dir” is the path to the folder where the Estimator should store and load checkpoints and event files. The “model_fn” parameter is a function that consumes the features, labels, mode and params in the following order:
def model_fn(features, labels, mode, params):
The Estimator will always supply those parameters when it executes the model function for training, evaluation or prediction. The features parameter contains a dictionary of tensors with the features that you want to feed to your network, the labels parameter contains a dictionary of tensors with the labels you want to use for training. These two parameters are generated by the input fn which will be explained later. The third parameter mode describes whether the “model_fn” is being called for training, evaluation or prediction. The final parameter params is a simple dictionary that can contain python variables and such that can be used during network definition (think total steps for learning rate schedule etc.).
Now that the Estimator object is initialized it can be used to start training and evaluating the network. Below is an excerpt of train.py which implements the above instructions to create an Estimator object after which it starts training and evaluating.
This snippet of code first sets up the config dictionary for the params, and uses it together with the “model_fn” to construct an Estimator Object. It then creates a “TrainSpec” with an “input_fn” and the “max_steps” the model should train for. A similar thing is done to create the “EvalSpec” where the “steps” are the number of steps of each evaluation, and “throttle_secs” defines the interval in seconds between each evaluation. The “tf.estimator.train_and_evaluate” is used to start the training and evaluation roster using the Estimator object, “TrainSpec” and “EvalSpec”. Finally an evaluation is explicitly called once more after the training finishes.
In my opinion the best things about Estimators, is that you can combine them easily with the Dataset class. Before the Estimator and Dataset class combination it was hard to prefetch and process examples on the CPU asynchronously from the GPU. Prefetching and processing on the CPU would in theory make sure there would be a batch of examples ready in memory any time the GPU was done processing the previous batch, in practice this was easier said then done. The problem with just assigning the fetch and process step to CPU is that unless it is done in parallel with the model processing on the GPU, the GPU still has to wait for the CPU to fetch the data from the storage disk and process it before it can start processing the next batch. For a long time queuerunners, asynchronous prefetching using python threads and other solutions were suggested, and in my experience none of them ever worked flawlessly and efficiently.
But using the Estimator class and combining it with the Dataset class is very easy, clean and works in parallel with the GPU. It allows the CPU to fetch, preprocess and que batches of examples such that there is always a new batch ready for the GPU. Using this method I have seen the utilization of a GPU stay close to 100% during training and global steps per second for small models(<10MB) increase 4 fold.
The Tensorflow Dataset class is designed as an E.T.L. process, which stands for Extract, Transform and Load. These steps will be defined soon, but this guide will only explain how to use tfrecords in combination with the Dataset class. For other formats (csv, numpy etc.) this page has a good write-up, however I suggest using tfrecords as they offer better performance and are easier to integrate with a Tensorflow development pipeline.
The whole E.T.L. process can be implemented using the Dataset class in only 7 lines of code as shown below. It might look complicated at first but read on for a detailed explanation of each lines functionality.
Extract
The first step in a Dataset input pipeline is to load the data from the tfrecords into memory. This starts with making a list of tfrecords available using a glob pattern e.g. “./Datasets/train-*.tfrecords” and the list_files function of the Dataset class. The parallel_interleave function is applied to the list of files, which ensures parallel extraction of the data as explained here. Finally a merged shuffle and repeat function is used to prefetch a certain number of examples from the tfrecords and shuffle them. The repeat ensures that there are always examples available by repeating from the start once the last example of every tfrecord is read.
files = tf.data.Dataset.list_files(glob_pattern, shuffle=True)dataset = files.apply(tf.contrib.data.parallel_interleave(                          lambda filename:                          tf.data.TFRecordDataset(filename),                          cycle_length=threads*2)                      )dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat                         (32*self.batch_size))
Transform
Now that the data is available in memory the next step is to transform it, preferably into something that does not need any further processing in order to be fed to the neural network input. A call to the dataset’s map function is required to do this as shown below, where “map_func” is the function applied to every individual example on the CPU and “num_parallel_calls” the number of parallel invocations of the “map_func” to use.
threads = multiprocessing.cpu_count()dataset = dataset.map(map_func=lambda example:                           _parse_function(example, self.image_size,                                self.num_classes,training=training),                            num_parallel_calls=threads)
In this case the “map_func” is the parse function shown below, this function processes an example from the tfrecords (created using this repo) and outputs a tuple of dictionaries containing tensors representing the features and labels respectively. Note the use of a lambda function to pass python variables separately from the example to the parse function, as the example is unparsed data from the tfrecord and provided by the Dataset class.
Keep in mind that this parse function only processes one example at the time as show by the tf.parse_single_example call, but does so a number of times in parallel. To prevent running into any CPU bottlenecks it is important to keep the parse function fast, some tips on how to do this can be found here. All the individually processed examples are then batched and ready for processing.
dataset = dataset.batch(batch_size=self.batch_size)
Load
The final step of the ETL process is loading the batched examples onto the accelerator (GPU) ready for processing. In the Dataset class this is achieved by prefetching, which is done by calling the prefetch function of the dataset.
dataset = dataset.prefetch(buffer_size=self.batch_size)
Prefetching uncouples the producer (Dataset object on CPU) from the consumer (GPU), this allows them to run in parallel for increased throughput.
Input Function
Once the whole E.T.L. process is fully defined and implemented, the “input_fn” can be created by initializing the iterator and grabbing the next example using the following line:
input_fn = dataset.make_one_shot_iterator().get_next()
This input function is used by the Estimator as an input for the model function.
A quick reminder, the model function the estimator invokes during training, evaluation and prediction, should accept the following arguments as explained earlier:
def model_fn(features, labels, mode, params):
The features and labels are in this case supplied by the Dataset class, the params are mostly for hyper parameters used during network initialization, but the mode (of type “tf.estimator.ModeKeys”) dictates what action the model is going to be performing. Each mode is used to setup the model for that specific purpose, the available modes are prediction, training and evaluation.
The different modes can be used by calling the respective functions (shown above) of an Estimator Object namely; “predict”, “train” and “evaluate”. The code-path of every mode has to return an “EstimatorSpec” with the required fields for that mode, e.g. when the mode is predict, it has to return an “EstimatorSpec” that includes the predictions field:
return tf.estimator.EstimatorSpec(mode, predictions=predictions)
Predict
The most basic mode is the prediction mode “tf.estimator.ModeKeys.PREDICT”, which as the name suggests is used to do predictions on data using the Estimator object. In this mode the “EstimatorSpec” expects a dictionary of tensors which will be executed and the results of which will be made available as numpy values to python.
In this excerpt you can see the predictions dictionary setup to generate classical image classification results. The if statement ensures that this code path is only executed when the predict function of an Estimator object is executed. The dictionary of tensors is passed to the “EstimatorSpec” as the predictions argument together with the mode. It is smart to define the prediction code-path first as it is the simplest, and since most of the code is used for training and evaluation as-well it can show problems early on.
Train
To train a model in the “tf.estimator.ModeKeys.TRAIN” mode it is necessary to create a so called “train_op”, this op is a tensor that when executed performs the back propagation to update the model. Simply put it is the minimize function of an optimizer such as the AdamOptimizer. The “train_op” and the scalar loss tensor are the minimum required arguments to create an “EstimatorSpec” for training. Below you can see an example of this being done.
Note the non-required arguments “training_hooks” and “scaffold”, these will be further explained later but in short they are used to add functionality to the setup and tear-down of a model and training session.
Evaluate
The final mode that needs a code-path in the model function is “tf.estimator.ModeKeys.EVAL”. The most important thing in order to perform an eval is the the metrics dictionary, this should be structured as a dictionary of tuples, where the first element of the tuple is a tensor containing the actual metric value and the second element is the tensor that updates the metric value. The update operation is necessary to ensure a reliable metric calculation over the whole validation set. Since it will often be impossible to evaluate the whole validation set in one batch, multiple batches have to be used. To prevent noise in the metric value due to per batch differences, the update operation is used to keep a running average (or gather all results) over all batches. This setup ensures the metric value is calculated over the whole validation set and not a single batch.
In the example above only the “loss” and “eval_metric_ops” are required arguments, the third argument “evaluation_hooks” is used to execute “tf.summary” operations as they are not automatically executed when running the Estimators evaluate function. In this example the “evaluation_hooks” are used to store images from the validation set to display using a Tensorboard. To achieve this a “SummarySaverHook” with the same output directory as the “model_dir” is initialized with a “tf.summary.image” operation and passed (encapsulated in an iterable) to the “EstimatorSpec”.
Now that I explained the advantages of using the Estimator class and how to use it, I hope you are excited to start using Estimators for your Tensorflow projects. However switching to Estimators from an existing code-base is not necessarily straightforward. As the Estimator class abstracts away most of the initialization and execution during training, any custom initialization and execution loops can no longer be an implementation using “tf.Session” and “sess.run()”. This could be a reason for someone with a large code-base not to transition to the Estimator class, as there is no easy and straightforward transition path. While this article is not a transition guide, I will clarify some of the new procedures for initialization and execution. This will hopefully fill in some of the gaps left by the official tutorials.
The main tools to influence the initialization and execution loop are the “Scaffold” object and the “SessionRunHook” object. The “Scaffold” is used for custom first time initialization of a model and can only be used to construct an “EstimatorSpec” in training mode. The “SessionRunHook” on the other hand can be used to construct an “EstimatorSpec” for each execution mode and is used each time train, evaluate or predict is called. Both the “Scaffold” and “SessionRunHook” provide certain functions that the Estimator class calls during use. Below you can see a timeline showing which function is called and when during the initialization and training process. This also shows that the Estimator under the hood still uses “tf.Session” and “sess.run”.
Scaffolds
A scaffold can be passed as the scaffold argument when constructing the “EstimatorSpec” for training. In a scaffold you can specify a number of different operations to be called during various stages of initialization. However for now we will focus on the “init_fn” argument as the others are more for distributed settings which this series will not cover. The “init_fn” is called after graph finalization but before the “sess.run” is called for the first time, it is only used once and thus perfect for custom variable initialization. A good example of this is if you want to finetune using a pretrained network and want to be selective about which variables to restore.
Above you can see an example of this implemented, the function checks for the existence of the fine-tune checkpoint and creates an initializer function using slim that filters on the “scope_name” variable. This function is then encapsulated in the function format that the Estimator expects which consumes the “Scaffold” object itself and a “tf.Session” object. Inside the “init_fn” the session is used to run the “initializer_fn”.
# setup fine tune scaffoldscaffold = tf.train.Scaffold(init_op=None,     init_fn=tools.fine_tune.init_weights(params["fine_tune_ckpt"]))# create estimator training specreturn tf.estimator.EstimatorSpec(tf.estimator.ModeKeys.TRAIN,                                loss=loss,                                train_op=train_op,scaffold=scaffold)
This “init_fn” can then be passed as an argument to construct a “Scaffold” object, as show above. This “Scaffold” object is used to construct the “EstimatorSpec” for the train mode. While this is a relatively simple example (and can also be achieved with “WarmStartSettings”), it can be easily expanded for more advanced onetime initialization functionality.
SessionRunHooks
As shown in the timeline earlier the “SessionRunHook” class can be used to alter certain parts of the training, evaluation or prediction loop. This is done using a concept known as hooking, I will not go into to much detail on what hooking is exactly as in this case it is pretty self explanatory. But using a “SessionRunHook” is slightly different from using a scaffold. Instead of initializing the “SessionRunHook” object with one or more functions as arguments it is used as a super class. It is then possible to overwrite the methods “begin”, “after_create_session”, “before_run”, “after_run” and “end”, to extend functionality. Below an excerpt showing how to overwrite the “begin” function is shown.
Keep in mind that the “begin” function does not get any arguments but each hook is slightly different, please refer here to check what arguments are passed to which function. Now one can create a “SessionRunHook” with the extended begin method by calling:
stats_hook = tools.stats.ModelStats("squeezenext",                                    params["model_dir"],                                    self.batch_size)
This object can then be used when constructing an “EstimatorSpec” for training, evaluating and predicting by passing it to the respective arguments “training_hooks”, “evaluation_hooks” and “prediction_hooks”. Note that each arguments ends with hooks, and they expect an iterable of one or multiple hooks so always encapsulate your hooks in an iterable.
Training and evaluation have extensively been covered, but prediction has not yet been fully explained. This is mostly because prediction is the easiest of the the 3 (train,eval,predict) Estimator modes, but there are still some pitfalls one can run into. I explained earlier how to setup an “EstimatorSpec” for prediction, but not how to use it. Below is a small example showing how to use an Estimator for prediction.
For prediction it does not make sense to make/use tfrecords, so instead the “numpy_input_fn” is used. To the “x” argument a dictionary of numpy arrays is passed containing the features, in this case an image.
Note: Depending on whether or not you are preprocess images in the “input_fn” during training, you would need to perform the same preprocessing here, either in numpy before passing to the “numpy_input_fn” or in Tensorflow.
With the numpy input function setup, a simple call to predict is enough to create a prediction. The “predictions” variable in the example above will not actually contain the result yet, instead it is an object of the generator class, to get the actual result you need to iterate over the it. This iteration will start the Tensorflow execution and produce the actual result.
I hope I was able to demystify some of the not so documented features of tensorflow Estimators and fill in the gaps left by the official tutorials. All the code in this post comes from this github repo, so please head over there for an example of Estimators used in combination with the Datasets class. If you have any more questions please leave them here, I’ll try my best to answer them.