# Migrate from Estimator to Keras APIs

<br />

|---|---|---|---|
| [![](https://www.tensorflow.org/images/tf_logo_32px.png) View on TensorFlow.org](https://www.tensorflow.org/guide/migrate/migrating_estimator) | [![](https://www.tensorflow.org/images/colab_logo_32px.png) Run in Google Colab](https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb) | [![](https://www.tensorflow.org/images/GitHub-Mark-32px.png) View source on GitHub](https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb) | [![](https://www.tensorflow.org/images/download_logo_32px.png)Download notebook](https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/migrating_estimator.ipynb) |

This guide demonstrates how to migrate from TensorFlow 1's `tf.estimator.Estimator` APIs to TensorFlow 2's [`tf.keras`](https://www.tensorflow.org/api_docs/python/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`](https://www.tensorflow.org/api_docs/python/tf/keras) APIs. You will also learn how to customize the training step by subclassing [`tf.keras.Model`](https://www.tensorflow.org/api_docs/python/tf/keras/Model) and using [`tf.GradientTape`](https://www.tensorflow.org/api_docs/python/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](https://www.tensorflow.org/guide/keras/custom_layers_and_models), gradient application, [training](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit), evaluation, and prediction.

(For migrating model/checkpoint saving workflows to TensorFlow 2, check out the [SavedModel](https://www.tensorflow.org/guide/migrate/saved_model) and [Checkpoint](https://www.tensorflow.org/guide/migrate/checkpoint_saved) migration guides.)

## Setup

Start with imports and a simple dataset:

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

    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:

    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:

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

Evaluate the program with the evaluation set:

    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`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) and [`Model.evaluate`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#evaluate) in TensorFlow 2. (You can learn more in the [Training and evaluation with the built-in methods](https://www.tensorflow.org/guide/keras/train_and_evaluate) guide.)

- Start by preparing the dataset pipeline with the [`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset) APIs.
- Define a simple Keras [Sequential](https://www.tensorflow.org/guide/keras/sequential_model) model with one linear ([`tf.keras.layers.Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense)) layer.
- Instantiate an Adagrad optimizer ([`tf.keras.optimizers.Adagrad`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad)).
- Configure the model for training by passing the `optimizer` variable and the mean-squared error (`"mse"`) loss to `Model.compile`.

    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`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit):

    model.fit(dataset)

Finally, evaluate the model with [`Model.evaluate`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#evaluate):

    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`](https://www.tensorflow.org/api_docs/python/tf/GradientTape) to perform forward and backward passes, while still taking advantage of the built-in training support, such as [`tf.keras.callbacks.Callback`](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/Callback) and [`tf.distribute.Strategy`](https://www.tensorflow.org/api_docs/python/tf/distribute/Strategy). (Learn more in [Customizing what happens in Model.fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit) and [Writing custom training loops from scratch](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch).)

In this example, start by creating a custom [`tf.keras.Model`](https://www.tensorflow.org/api_docs/python/tf/keras/Model) by subclassing [`tf.keras.Sequential`](https://www.tensorflow.org/api_docs/python/tf/keras/Sequential) that overrides [`Model.train_step`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#train_step). (Learn more about [subclassing tf.keras.Model](https://www.tensorflow.org/guide/keras/custom_layers_and_models)). 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.

    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`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset).
- Define a simple model with one [`tf.keras.layers.Dense`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Dense) layer.
- Instantiate Adagrad ([`tf.keras.optimizers.Adagrad`](https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad))
- Configure the model for training with [`Model.compile`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile), while using mean-squared error (`"mse"`) as the loss function.

    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`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit) to train the model:

    model.fit(dataset)

And, finally, evaluate the program with [`Model.evaluate`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#evaluate):

    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](https://www.tensorflow.org/guide/keras/train_and_evaluate)
- Guide: [Customize what happens in Model.fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit)
- Guide: [Writing a training loop from scratch](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch)
- Guide: [Making new Keras layers and models via subclassing](https://www.tensorflow.org/guide/keras/custom_layers_and_models)

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

- [Migrate from TPUEstimator to TPUStrategy](https://www.tensorflow.org/guide/migrate/tpu_estimator)
- [Migrate single-worker multiple-GPU training](https://www.tensorflow.org/guide/migrate/mirrored_strategy)
- [Migrate multi-worker CPU/GPU training](https://www.tensorflow.org/guide/migrate/multi_worker_cpu_gpu_training)