Notebook

Copyright 2021 The TensorFlow Authors.¶

In [ ]:

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Migrate from Estimator to Keras APIs¶

View on TensorFlow.org

Run in Google Colab

View source on GitHub

Download notebook

This guide demonstrates how to migrate from TensorFlow 1's tf.estimator.Estimator APIs to TensorFlow 2's tf.keras APIs. First, you will set up and run a basic model for training and evaluation with tf.estimator.Estimator. Then, you will perform the equivalent steps in TensorFlow 2 with the tf.keras APIs. You will also learn how to customize the training step by subclassing tf.keras.Model and using tf.GradientTape.

In TensorFlow 1, the high-level tf.estimator.Estimator APIs let you train and evaluate a model, as well as perform inference and save your model (for serving).
In TensorFlow 2, use the Keras APIs to perform the aforementioned tasks, such as model building, gradient application, training, evaluation, and prediction.

(For migrating model/checkpoint saving workflows to TensorFlow 2, check out the SavedModel and Checkpoint migration guides.)

Setup¶

Start with imports and a simple dataset:

In [ ]:

import tensorflow as tf
import tensorflow.compat.v1 as tf1

In [ ]:

features = [[1., 1.5], [2., 2.5], [3., 3.5]]
labels = [[0.3], [0.5], [0.7]]
eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]
eval_labels = [[0.8], [0.9], [1.]]

TensorFlow 1: Train and evaluate with tf.estimator.Estimator¶

This example shows how to perform training and evaluation with tf.estimator.Estimator in TensorFlow 1.

Start by defining a few functions: an input function for the training data, an evaluation input function for the evaluation data, and a model function that tells the Estimator how the training op is defined with the features and labels:

In [ ]:

def _input_fn():
  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)

def _eval_input_fn():
  return tf1.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

def _model_fn(features, labels, mode):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

Instantiate your Estimator, and train the model:

In [ ]:

estimator = tf1.estimator.Estimator(model_fn=_model_fn)
estimator.train(_input_fn)

Evaluate the program with the evaluation set:

In [ ]:

estimator.evaluate(_eval_input_fn)

TensorFlow 2: Train and evaluate with the built-in Keras methods¶

This example demonstrates how to perform training and evaluation with Keras Model.fit and Model.evaluate in TensorFlow 2. (You can learn more in the Training and evaluation with the built-in methods guide.)

Start by preparing the dataset pipeline with the tf.data.Dataset APIs.
Define a simple Keras Sequential model with one linear (tf.keras.layers.Dense) layer.
Instantiate an Adagrad optimizer (tf.keras.optimizers.Adagrad).
Configure the model for training by passing the optimizer variable and the mean-squared error ("mse") loss to Model.compile.

In [ ]:

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

With that, you are ready to train the model by calling Model.fit:

In [ ]:

model.fit(dataset)

Finally, evaluate the model with Model.evaluate:

In [ ]:

model.evaluate(eval_dataset, return_dict=True)

TensorFlow 2: Train and evaluate with a custom training step and built-in Keras methods¶

In TensorFlow 2, you can also write your own custom training step function with tf.GradientTape to perform forward and backward passes, while still taking advantage of the built-in training support, such as tf.keras.callbacks.Callback and tf.distribute.Strategy. (Learn more in Customizing what happens in Model.fit and Writing custom training loops from scratch.)

In this example, start by creating a custom tf.keras.Model by subclassing tf.keras.Sequential that overrides Model.train_step. (Learn more about subclassing tf.keras.Model). Inside that class, define a custom train_step function that for each batch of data performs a forward pass and backward pass during one training step.

In [ ]:

class CustomModel(tf.keras.Sequential):
  """A custom sequential model that overrides `Model.train_step`."""

  def train_step(self, data):
    batch_data, labels = data

    with tf.GradientTape() as tape:
      predictions = self(batch_data, training=True)
      # Compute the loss value (the loss function is configured
      # in `Model.compile`).
      loss = self.compiled_loss(labels, predictions)

    # Compute the gradients of the parameters with respect to the loss.
    gradients = tape.gradient(loss, self.trainable_variables)
    # Perform gradient descent by updating the weights/parameters.
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    # Update the metrics (includes the metric that tracks the loss).
    self.compiled_metrics.update_state(labels, predictions)
    # Return a dict mapping metric names to the current values.
    return {m.name: m.result() for m in self.metrics}

Next, as before:

Prepare the dataset pipeline with tf.data.Dataset.
Define a simple model with one tf.keras.layers.Dense layer.
Instantiate Adagrad (tf.keras.optimizers.Adagrad)
Configure the model for training with Model.compile, while using mean-squared error ("mse") as the loss function.

In [ ]:

dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)
eval_dataset = tf.data.Dataset.from_tensor_slices(
      (eval_features, eval_labels)).batch(1)

model = CustomModel([tf.keras.layers.Dense(1)])
optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)

model.compile(optimizer=optimizer, loss="mse")

Call Model.fit to train the model:

In [ ]:

model.fit(dataset)

And, finally, evaluate the program with Model.evaluate:

In [ ]:

model.evaluate(eval_dataset, return_dict=True)

Next steps¶

Additional Keras resources you may find useful:

Guide: Training and evaluation with the built-in methods
Guide: Customize what happens in Model.fit
Guide: Writing a training loop from scratch
Guide: Making new Keras layers and models via subclassing

The following guides can assist with migrating distribution strategy workflows from tf.estimator APIs: