How to move from tf.contrib.learn Estimator to core Tensorflow tf.Estimator | by Lak Lakshmanan | Towards Data Science

I have a bunch of machine learning models built using the “original” Estimator API, the one in tf.contrib.learn. In Tensorflow 1.2, parts of it moved to core and in Tensorflow 1.4, the remaining pieces I need finally arrived in core Tensorflow. So, time for me to begin the process of moving to the core Estimator API.
When the Tensorflow team moves code from contrib to core, they take the opportunity to clean up the APIs. The alternative would be to do expensive code reviews of everything before it is accepted into contrib, and not just of the things that prove useful. That would slow down innovation ...
In any case, if we found an experimental package in contrib to be useful, we should be prepared to spend a bit of time and effort to migrate our code once it moves to core. This will usually not be just a matter of changing the package name (for example) from tf.contrib.learn.Estimator to tf.Estimator and calling it a day.
Let’s look at the process of moving from tf.contrib.learn to tf.estimator in four parts:
If you are moving a canned estimator (e.g., tf.contrib.learn.DNNClassifier to tf.Estimator.DNNClassifier)If you are moving a custom estimator (e.g., if you did something similar to my blog posts on text convolution or time-series forecasting and so need to change from tf.contrib.Estimator(model_fn) to tf.Estimator(model_fn).How to remove tf.contrib.learn.Experiment and learn_runner from your code.Changes to your Cloud ML Engine scripts
If you are moving a canned estimator (e.g., tf.contrib.learn.DNNClassifier to tf.Estimator.DNNClassifier)
If you are moving a custom estimator (e.g., if you did something similar to my blog posts on text convolution or time-series forecasting and so need to change from tf.contrib.Estimator(model_fn) to tf.Estimator(model_fn).
How to remove tf.contrib.learn.Experiment and learn_runner from your code.
Changes to your Cloud ML Engine scripts
Woah! That’s a lot of changes! Considering the sheer number of changes here, how likely are you to use classes from contrib in the future? Are you thinking that you would rather wait for the dust to settle, and use things after they make their way into core?
Well, sure, if you want a stable API, you would stick with just the things in core Tensorflow. But I think you are like me and believe Bernard Shaw’s maxim that progress is impossible without change, and you decided to use contrib knowing that you would have to make some changes in the future.
The jig’s up, though, and it’s time to pay the piper! Let’s walk through the changes one-by-one.
The changes to canned estimator primarily involve the feature columns. Take these lines:
import tensorflow.contrib.layers as layersis_male = layers.sparse_column_with_keys('is_male',                             keys=['True', 'False'])age = layers.real_valued_column('mother_age')b_age = layers.bucketized_column(age, np.arange(15,45,1).tolist())crossed = layers.crossed_column([b_age, ...], 200)embed = layers.embedding_column(crossed, 3)
The feature column package and method names have changed. Here’s what the replacements look like:
import tensorflow.feature_column as fcis_male = fc.categorical_column_with_vocabulary_list('is_male',                                                    ['True', 'False'])age = fc.numeric_column('mother_age')b_age = fc.bucketized_column(age,                        boundaries=np.arange(15,45,1).tolist())crossed = fc.crossed_column([b_age, ...],                         hash_bucket_size=200)embed = fc.embedding_column(crossed, 3)
The Tensorflow documentation on feature columns is quite helpful in case some of these terms are unfamiliar. The changed method names (categorical instead of sparse and numeric instead of real) get to the characteristic of the fields rather than their representation, and the new APIs have more flexibility.
Once you have the feature columns, creating the estimator remains pretty much the same. Change, for example,
tf.contrib.learn.DNNLinearCombinedRegressor
to:
tf.estimator.DNNLinearCombinedRegressor
and you are done.
However, you should also tighten up your serving input_fn (this is the function that will be used by Tensorflow Serving to parse the received JSON, for example, and feed it into your model). You probably had something like:
return tf.contrib.learn.utils.input_fn_utils.InputFnOps(     features,     None,     feature_placeholders)
Now, it will become:
return tf.estimator.export.ServingInputReceiver(           features,            feature_placeholders)
Essentially that second arg, that nobody ever needed, has gone to its resting place.
In the process of moving to core, the custom estimator has been made much more flexible and powerful.
In your custom estimator, you used to have a model function that was defined like this:
def cnn_model(features, target, mode):   # stuff ...   return tf.contrib.learn.ModelFnOps(      mode=mode,      predictions=predictions_dict,      loss=loss,      train_op=train_op)
This now becomes:
def cnn_model(features, labels, mode, params):    # STUFF    if mode == tf.estimator.ModeKeys.TRAIN or mode == tf.estimator.ModeKeys.EVAL:       eval_metrics = {'acc': tf.metrics.accuracy(tf.argmax(logits,1), labels)}    else:       eval_metrics = None      return tf.estimator.EstimatorSpec(      mode=mode,      predictions=predictions_dict,      loss=loss,      train_op=train_op,      eval_metric_ops=eval_metrics,      export_outputs={'classes': tf.estimator.export.PredictOutput(predictions_dict)}    )
Let’s unpack that. These are the key changes:
Note that you now get an extra parameter. You can use this to pass in an extra dictionary to your model. This comes in quite handy when you are using Cloud ML Engine and you want to hyperparameter tune the model parameters. More on this shortly.The mode that used to be something like tf.contrib.learn.ModeKeys.TRAIN is now tf.estimator.ModeKeys.TRAIN. No big deal here. The numbers they map to are still the same, so even if you don’t change your code, nothing crazy will happen. But better to clean it up and use the core package.Instead of returning a tf.contrib.learn.ModelFnOps you now return an tf.estimator.EstimatorSpec, but EstimatorSpec has a few new bells and whistles. Specifically, you can now specify a list of metrics to evaluate. EstimatorSpec also provides the ability to export different slices of your predictions dictionary.
Note that you now get an extra parameter. You can use this to pass in an extra dictionary to your model. This comes in quite handy when you are using Cloud ML Engine and you want to hyperparameter tune the model parameters. More on this shortly.
The mode that used to be something like tf.contrib.learn.ModeKeys.TRAIN is now tf.estimator.ModeKeys.TRAIN. No big deal here. The numbers they map to are still the same, so even if you don’t change your code, nothing crazy will happen. But better to clean it up and use the core package.
Instead of returning a tf.contrib.learn.ModelFnOps you now return an tf.estimator.EstimatorSpec, but EstimatorSpec has a few new bells and whistles. Specifically, you can now specify a list of metrics to evaluate. EstimatorSpec also provides the ability to export different slices of your predictions dictionary.
The model function itself was passed into an estimator like this:
tf.contrib.learn.Estimator(model_fn=cnn_model,            model_dir=output_dir),
Instead, you will pass it like this:
training_config=tf.estimator.RunConfig(save_checkpoints_steps=1000)hparams = parser.parse_args().__dict__tf.estimator.Estimator(model_fn=cnn_model,                                      model_dir=output_dir,                                   config=training_config,                                   params=hparams),
Notice that I’m passing in the command-line parameters as params. This is what becomes the fourth parameter to cnn_model() above — in Cloud ML Engine, this pass-through allows me to hyper-parameter tune the model.
Experiment, pardon the pun, was a failed experiment. The replacement core Tensorflow functionality is to use tf.Estimator.train_and_evaluate().
You probably had a setup like this:
def make_experiment_fn(output_dir):  def experiment_fn(output_dir):     wide, deep = get_wide_deep()     return tf.contrib.learn.Experiment(          estimator,            train_input_fn=read_dataset('train'),        eval_input_fn=read_dataset('eval'),        eval_metrics={            'rmse': tf.contrib.learn.MetricSpec(                metric_fn=metrics.streaming_root_mean_squared_error            )        },        train_input_fn=read_dataset('train', PATTERN),        eval_input_fn=read_dataset('eval', PATTERN),        export_strategies=[saved_model_export_utils.make_export_strategy(             serving_input_fn,             default_output_alternative_key=None,             exports_to_keep=1         )],        train_steps=TRAIN_STEPS     )   return experiment_fnlearn_runner.run(make_experiment_fn, output_dir)
This will now become:
train_spec=tf.estimator.TrainSpec(                         input_fn=read_dataset('train', PATTERN),                         max_steps=TRAIN_STEPS)exporter = tf.estimator.LatestExporter('exporter',serving_input_fn)eval_spec=tf.estimator.EvalSpec(                         input_fn=read_dataset('eval', PATTERN),                         steps=None,                         exporters=exporter)tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
Essentially, where you used to call learn_runner.run(), you now call tf.estimator.train_and_evaluate(). You pass it the estimator and train_spec which encapsulates details about your training data and eval_spec which does the same for your evaluation data. The exporting happens during evaluation, so you specify your exporter (either LatestExporter or FinalExporter) there.
Note that in Experiment, you used to specify eval_metrics, but that has moved into the EstimatorSpec (see previous section, on custom estimator). It’s not clear to me how you could get extra metrics evaluated if you are using a canned estimator ... Fortunately, I haven’t run across that need yet.
What changes do you have to make to your development environment?
Well, Cloud Datalab now comes with Tensorflow 1.4, so just launch a new instance and run your notebooks in that. No other changes needed.
In your gcloud commands where you submit your training jobs, specify a runtime_version of 1.4 instead of 1.2. So change:
gcloud ml-engine jobs submit training $JOBNAME \   ...   --runtime-version 1.2 \   ...
to:
gcloud ml-engine jobs submit training $JOBNAME \   ...   --runtime-version 1.4 \   ...
The saved_model_export_utils in contrib used to write the exported models to a directory called Servo. Well, Servo was the internal project codename for the Google project that built Tensorflow Serving ... The directory is now called “exporter” instead of “Servo”. Presumably, this is in case another team comes along and builds something better. So if you had a script like:
MODEL_LOCATION=$(gsutil ls gs://${BUCKET}/babyweight/trained_model/export/Servo/ | tail -1)
you will change it to:
MODEL_LOCATION=$(gsutil ls gs://${BUCKET}/babyweight/trained_model/export/exporter/ | tail -1)
And when you deploy the model, make sure to specify the Tensorflow runtime version:
gcloud ml-engine versions create ${MODEL_VERSION} --model ${MODEL_NAME} --origin ${MODEL_LOCATION} --runtime-version 1.4
Everything else remains the same.