{ "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", "\n", " \n", " \n", " \n", " \n", "
\n", " \n", " \n", " View on TensorFlow.org\n", " \n", " \n", " \n", " Run in Google Colab\n", " \n", " \n", " \n", " View source on GitHub\n", " \n", " Download notebook\n", "
" ] }, { "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 }