{ "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 single-worker multiple-GPU training\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 the single-worker multiple-GPU workflows from TensorFlow 1 to TensorFlow 2.\n", "\n", "To perform synchronous training across multiple GPUs on one machine:\n", "\n", "- In TensorFlow 1, you use the `tf.estimator.Estimator` APIs with `tf.distribute.MirroredStrategy`.\n", "- In TensorFlow 2, you can use [Keras Model.fit](https://www.tensorflow.org/tutorials/distribute/keras) or [a custom training loop](https://www.tensorflow.org/tutorials/distribute/custom_training) with `tf.distribute.MirroredStrategy`. Learn more in the [Distributed training with TensorFlow](https://www.tensorflow.org/guide/distributed_training#mirroredstrategy) guide." ] }, { "cell_type": "markdown", "metadata": { "id": "YdZSoIXEbhg-" }, "source": [ "## Setup" ] }, { "cell_type": "markdown", "metadata": { "id": "6d466b39d0db" }, "source": [ "Start with imports and a simple dataset for demonstration purposes:" ] }, { "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: Single-worker distributed training with tf.estimator.Estimator" ] }, { "cell_type": "markdown", "metadata": { "id": "A9560BqEOTpb" }, "source": [ "This example demonstrates the TensorFlow 1 canonical workflow of single-worker multiple-GPU training. You need to set the distribution strategy (`tf.distribute.MirroredStrategy`) through the `config` parameter of the `tf.estimator.Estimator`:" ] }, { "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)\n", "\n", "strategy = tf1.distribute.MirroredStrategy()\n", "config = tf1.estimator.RunConfig(\n", " train_distribute=strategy, eval_distribute=strategy)\n", "estimator = tf1.estimator.Estimator(model_fn=_model_fn, config=config)\n", "\n", "train_spec = tf1.estimator.TrainSpec(input_fn=_input_fn)\n", "eval_spec = tf1.estimator.EvalSpec(input_fn=_eval_input_fn)\n", "tf1.estimator.train_and_evaluate(estimator, train_spec, eval_spec)" ] }, { "cell_type": "markdown", "metadata": { "id": "KEmzBjfnsxwT" }, "source": [ "## TensorFlow 2: Single-worker training with Keras" ] }, { "cell_type": "markdown", "metadata": { "id": "fkgkGf_AOaRR" }, "source": [ "When migrating to TensorFlow 2, you can use the Keras APIs with `tf.distribute.MirroredStrategy`.\n", "\n", "If you use the `tf.keras` APIs for model building and Keras `Model.fit` for training, the main difference is instantiating the Keras model, an optimizer, and metrics in the context of `Strategy.scope`, instead of defining a `config` for `tf.estimator.Estimator`.\n", "\n", "If you need to use a custom training loop, check out the [Using tf.distribute.Strategy with custom training loops](https://www.tensorflow.org/guide/distributed_training#using_tfdistributestrategy_with_custom_training_loops) guide." ] }, { "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)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Kip65sYBlKiu" }, "outputs": [], "source": [ "strategy = tf.distribute.MirroredStrategy()\n", "with strategy.scope():\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')\n", "model.fit(dataset)\n", "model.evaluate(eval_dataset, return_dict=True)" ] }, { "cell_type": "markdown", "metadata": { "id": "0431f3935485" }, "source": [ "## Next steps" ] }, { "cell_type": "markdown", "metadata": { "id": "a68d2a99f79b" }, "source": [ "To learn more about distributed training with `tf.distribute.MirroredStrategy` in TensorFlow 2, check out the following documentation:\n", "\n", "- The [Distributed training on one machine with Keras](../../tutorials/distribute/keras) tutorial\n", "- The [Distributed training on one machine with a custom training loop](../../tutorials/distribute/custom_training) tutorial\n", "- The [Distributed training with TensorFlow](../../guide/distributed_training) guide\n", "- The [Using multiple GPUs](../../guide/gpu#using_multiple_gpus) guide\n", "- The [Optimize the performance on the multi-GPU single host (with the TensorFlow Profiler)](../../guide/gpu_performance_analysis#2_optimize_the_performance_on_the_multi-gpu_single_host) guide" ] } ], "metadata": { "colab": { "collapsed_sections": [], "name": "mirrored_strategy.ipynb", "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }