{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wJcYs_ERTnnI"
      },
      "outputs": [],
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HMUDt0CiUJk9",
        "vscode": {
          "languageId": "python"
        }
      },
      "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": [
        "# 迁移单工作进程多 GPU 训练\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://tensorflow.google.cn/guide/migrate/mirrored_strategy\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\">在 TensorFlow.org 上查看</a>   </td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/migrate/mirrored_strategy.ipynb\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\">在 Google Colab 运行</a>\n",
        "</td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/migrate/mirrored_strategy.ipynb\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\">在 Github 上查看源代码</a>\n",
        "</td>\n",
        "  <td>     <a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/guide/migrate/mirrored_strategy.ipynb\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\">下载笔记本</a>   </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "meUTrR4I6m1C"
      },
      "source": [
        "本指南演示了如何将单工作进程多 GPU 工作流从 TensorFlow 1 迁移到 TensorFlow 2。\n",
        "\n",
        "要在一台机器上跨多个 GPU 执行同步训练，请执行以下操作：\n",
        "\n",
        "- 在 TensorFlow 1 中，将 `tf.estimator.Estimator` API 与 `tf.distribute.MirroredStrategy` 一起使用。\n",
        "- 在 TensorFlow 2 中，可以使用 [Keras Model.fit](https://tensorflow.google.cn/tutorials/distribute/keras) 或带有 `tf.distribute.MirroredStrategy` 的[自定义训练循环](https://tensorflow.google.cn/tutorials/distribute/custom_training)。有关详情，请参阅[使用 TensorFlow 进行分布式训练](https://tensorflow.google.cn/guide/distributed_training#mirroredstrategy)指南。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdZSoIXEbhg-"
      },
      "source": [
        "## 安装"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6d466b39d0db"
      },
      "source": [
        "从导入和用于演示目的的简单数据集开始："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iE0vSfMXumKI",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.compat.v1 as tf1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7rnGxsXtDkV",
        "vscode": {
          "languageId": "python"
        }
      },
      "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：使用 tf.estimator.Estimator 进行单工作进程分布式训练"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A9560BqEOTpb"
      },
      "source": [
        "此示例演示了单工作进程多 GPU 训练的 TensorFlow 1 规范工作流。您需要通过 `tf.estimator.Estimator` 的 `config` 参数设置分布策略 (`tf.distribute.MirroredStrategy`)："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lqe9obf7suIj",
        "vscode": {
          "languageId": "python"
        }
      },
      "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：使用 Keras 进行单工作进程训练"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fkgkGf_AOaRR"
      },
      "source": [
        "迁移到 TensorFlow 2 时，可以将 Keras API 与 `tf.distribute.MirroredStrategy` 一起使用。\n",
        "\n",
        "如果您使用 `tf.keras` API 进行模型构建，并使用 Keras `Model.fit` 进行训练，那么主要区别在于，这会在 `Strategy.scope` 的上下文中实例化 Keras 模型、优化器和指标，而不是为 `tf.estimator.Estimator` 定义 `config`。\n",
        "\n",
        "如果您需要使用自定义训练循环，请查看[将 tf.distribute.Strategy 与自定义训练循环一起使用](https://tensorflow.google.cn/guide/distributed_training#using_tfdistributestrategy_with_custom_training_loops)指南。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "atVciNgPs0fw",
        "vscode": {
          "languageId": "python"
        }
      },
      "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",
        "vscode": {
          "languageId": "python"
        }
      },
      "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": [
        "## 后续步骤"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a68d2a99f79b"
      },
      "source": [
        "要详细了解如何在 TensorFlow 2 中使用 `tf.distribute.MirroredStrategy` 进行分布式训练，请查看以下文档：\n",
        "\n",
        "- [使用 Keras 在一台机器上进行分布式训练](../../tutorials/distribute/keras)教程\n",
        "- [使用自定义训练循环在一台机器上进行分布式训练](../../tutorials/distribute/custom_training)教程\n",
        "- [使用 TensorFlow 进行分布式训练](../../guide/distributed_training)指南\n",
        "- [使用多个 GPU](../../guide/gpu#using_multiple_gpus) 指南\n",
        "- [优化多 GPU 单主机上的性能（使用 TensorFlow Profiler）](../../guide/gpu_performance_analysis#2_optimize_the_performance_on_the_multi-gpu_single_host)指南"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "mirrored_strategy.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}