site/en/guide/migrate/migrating_estimator.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "77z2OchJTk0l"
      },
      "source": [
        "# Migrate from Estimator to Keras APIs\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/migrate/migrating_estimator\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/migrate/migrating_estimator.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/migrate/migrating_estimator.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "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`.\n",
        "\n",
        "- 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).\n",
        "- 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.\n",
        "\n",
        "(For migrating model/checkpoint saving workflows to TensorFlow 2, check out the [SavedModel](saved_model.ipynb) and [Checkpoint](checkpoint_saved.ipynb) migration guides.)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Start with imports and a simple dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7rnGxsXtDkV"
      },
      "outputs": [],
      "source": [
        "features = [[1., 1.5], [2., 2.5], [3., 3.5]]\n",
        "labels = [[0.3], [0.5], [0.7]]\n",
        "eval_features = [[4., 4.5], [5., 5.5], [6., 6.5]]\n",
        "eval_labels = [[0.8], [0.9], [1.]]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4uXff1BEssdE"
      },
      "source": [
        "## TensorFlow 1: Train and evaluate with tf.estimator.Estimator\n",
        "\n",
        "This example shows how to perform training and evaluation with `tf.estimator.Estimator` in TensorFlow 1.\n",
        "\n",
        "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:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lqe9obf7suIj"
      },
      "outputs": [],
      "source": [
        "def _input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices((features, labels)).batch(1)\n",
        "\n",
        "def _eval_input_fn():\n",
        "  return tf1.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).batch(1)\n",
        "\n",
        "def _model_fn(features, labels, mode):\n",
        "  logits = tf1.layers.Dense(1)(features)\n",
        "  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)\n",
        "  optimizer = tf1.train.AdagradOptimizer(0.05)\n",
        "  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())\n",
        "  return tf1.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "44bf417bf9c0"
      },
      "source": [
        "Instantiate your `Estimator`, and train the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "922720812527"
      },
      "outputs": [],
      "source": [
        "estimator = tf1.estimator.Estimator(model_fn=_model_fn)\n",
        "estimator.train(_input_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "17c9933c2d89"
      },
      "source": [
        "Evaluate the program with the evaluation set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HsOpjW5plH9Q"
      },
      "outputs": [],
      "source": [
        "estimator.evaluate(_eval_input_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEmzBjfnsxwT"
      },
      "source": [
        "## TensorFlow 2: Train and evaluate with the built-in Keras methods\n",
        "\n",
        "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](https://www.tensorflow.org/guide/keras/train_and_evaluate) guide.)\n",
        "\n",
        "- Start by preparing the dataset pipeline with the `tf.data.Dataset` APIs.\n",
        "- Define a simple Keras [Sequential](https://www.tensorflow.org/guide/keras/sequential_model) model with one linear (`tf.keras.layers.Dense`) layer.\n",
        "- Instantiate an Adagrad optimizer (`tf.keras.optimizers.Adagrad`).\n",
        "- Configure the model for training by passing the `optimizer` variable and the mean-squared error (`\"mse\"`) loss to `Model.compile`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atVciNgPs0fw"
      },
      "outputs": [],
      "source": [
        "dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)\n",
        "eval_dataset = tf.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).batch(1)\n",
        "\n",
        "model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])\n",
        "optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
        "\n",
        "model.compile(optimizer=optimizer, loss=\"mse\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ed17a6291959"
      },
      "source": [
        "With that, you are ready to train the model by calling `Model.fit`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a0b732534501"
      },
      "outputs": [],
      "source": [
        "model.fit(dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74767288a2ea"
      },
      "source": [
        "Finally, evaluate the model with `Model.evaluate`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kip65sYBlKiu"
      },
      "outputs": [],
      "source": [
        "model.evaluate(eval_dataset, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuVYN0CHs5sD"
      },
      "source": [
        "## TensorFlow 2: Train and evaluate with a custom training step and built-in Keras methods"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gHx_RUL8xcJ3"
      },
      "source": [
        "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](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).)\n",
        "\n",
        "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](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.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rSz_y0zOs8h2"
      },
      "outputs": [],
      "source": [
        "class CustomModel(tf.keras.Sequential):\n",
        "  \"\"\"A custom sequential model that overrides `Model.train_step`.\"\"\"\n",
        "\n",
        "  def train_step(self, data):\n",
        "    batch_data, labels = data\n",
        "\n",
        "    with tf.GradientTape() as tape:\n",
        "      predictions = self(batch_data, training=True)\n",
        "      # Compute the loss value (the loss function is configured\n",
        "      # in `Model.compile`).\n",
        "      loss = self.compiled_loss(labels, predictions)\n",
        "\n",
        "    # Compute the gradients of the parameters with respect to the loss.\n",
        "    gradients = tape.gradient(loss, self.trainable_variables)\n",
        "    # Perform gradient descent by updating the weights/parameters.\n",
        "    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
        "    # Update the metrics (includes the metric that tracks the loss).\n",
        "    self.compiled_metrics.update_state(labels, predictions)\n",
        "    # Return a dict mapping metric names to the current values.\n",
        "    return {m.name: m.result() for m in self.metrics}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee7c4f94d69b"
      },
      "source": [
        "Next, as before:\n",
        "- Prepare the dataset pipeline with `tf.data.Dataset`.\n",
        "- Define a simple model with one `tf.keras.layers.Dense` layer.\n",
        "- Instantiate Adagrad (`tf.keras.optimizers.Adagrad`)\n",
        "- Configure the model for training with `Model.compile`, while using mean-squared error (`\"mse\"`) as the loss function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "01fcc2b1292c"
      },
      "outputs": [],
      "source": [
        "dataset = tf.data.Dataset.from_tensor_slices((features, labels)).batch(1)\n",
        "eval_dataset = tf.data.Dataset.from_tensor_slices(\n",
        "      (eval_features, eval_labels)).batch(1)\n",
        "\n",
        "model = CustomModel([tf.keras.layers.Dense(1)])\n",
        "optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)\n",
        "\n",
        "model.compile(optimizer=optimizer, loss=\"mse\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "844543802ff5"
      },
      "source": [
        "Call `Model.fit` to train the model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "211be3620765"
      },
      "outputs": [],
      "source": [
        "model.fit(dataset)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c93b9d6fc9d7"
      },
      "source": [
        "And, finally, evaluate the program with `Model.evaluate`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nYO2wI1SlNCG"
      },
      "outputs": [],
      "source": [
        "model.evaluate(eval_dataset, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e9b5c9a4747b"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "Additional Keras resources you may find useful:\n",
        "\n",
        "- Guide: [Training and evaluation with the built-in methods](https://www.tensorflow.org/guide/keras/train_and_evaluate)\n",
        "- Guide: [Customize what happens in Model.fit](https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit)\n",
        "- Guide: [Writing a training loop from scratch](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch)\n",
        "- Guide: [Making new Keras layers and models via subclassing](https://www.tensorflow.org/guide/keras/custom_layers_and_models)\n",
        "\n",
        "The following guides can assist with migrating distribution strategy workflows from `tf.estimator` APIs:\n",
        "\n",
        "- [Migrate from TPUEstimator to TPUStrategy](tpu_estimator.ipynb)\n",
        "- [Migrate single-worker multiple-GPU training](mirrored_strategy.ipynb)\n",
        "- [Migrate multi-worker CPU/GPU training](multi_worker_cpu_gpu_training.ipynb)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "migrating_estimator.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}