<commit_before>{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "multi_worker.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "Tce3stUlHN0L"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "metadata": {
        "colab_type": "text",
        "id": "Tce3stUlHN0L"
      },
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "tuOe1ymfHZPu",
        "colab": {}
      },
      "cell_type": "code",
      "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."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MfBg1C5NB3X0"
      },
      "cell_type": "markdown",
      "source": [
        "# Multi-worker Training in TensorFlow\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/alpha/tutorials/distribute/multi_worker\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />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/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "xHxb-dlhMIzW"
      },
      "cell_type": "markdown",
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how `tf.distribute.Strategy` can be used for distributed multi-worker training with `tf.estimator`.  If you write your code using `tf.estimator`, and you're interested in scaling beyond a single machine with high performance, this tutorial is for you.\n",
        "\n",
        "Before getting started, please read the [`tf.distribute.Strategy` guide](../../guide/distribute_strategy.ipynb).  The [multi-GPU training tutorial](./keras.ipynb) is also relevant, because this tutorial uses the same model.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MUXex9ctTuDB"
      },
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "First, setup TensorFlow and the necessary imports."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "IqR2PQG4ZaZ0",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import, division, print_function"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bnYxvfLD-LW-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!pip install tensorflow-gpu==2.0.0-alpha0\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import os, json"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "hPBuZUNSZmrQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Input Function\n",
        "\n",
        "This tutorial uses the MNIST dataset from [TensorFlow Datasets](https://www.tensorflow.org/datasets).  The code here is similar to the [multi-GPU training tutorial](./keras.ipynb) with one key difference: when using Estimator for multi-worker training, it is necessary to shard the dataset by the number of workers to ensure model convergence.  The input data is sharded by worker index, so that each worker processes 1/`num_workers` distinct portions of the dataset."
      ]
    },
    {
      "metadata": {
        "id": "dma_wUAxZqo2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def input_fn(mode, input_context=None):\n",
        "  datasets, ds_info = tfds.load(name='mnist',\n",
        "                                with_info=True,\n",
        "                                as_supervised=True)\n",
        "  mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else\n",
        "                   datasets['test'])\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255\n",
        "    return image, label\n",
        "\n",
        "  if input_context:\n",
        "    mnist_dataset = mnist_dataset.apply(tf.data.experimental.filter_for_shard(\n",
        "        input_context.num_input_pipelines, input_context.input_pipeline_id))\n",
        "  return mnist_dataset.map(scale).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
      ],
      "execution_count": 0,<commit_msg>Change `tf.data.experimental.filter_for_shard` to `tf.data.Dataset.shard`.<commit_after>{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "multi_worker.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "Tce3stUlHN0L"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "metadata": {
        "colab_type": "text",
        "id": "Tce3stUlHN0L"
      },
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "tuOe1ymfHZPu",
        "colab": {}
      },
      "cell_type": "code",
      "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."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MfBg1C5NB3X0"
      },
      "cell_type": "markdown",
      "source": [
        "# Multi-worker Training in TensorFlow\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/alpha/tutorials/distribute/multi_worker\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />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/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "xHxb-dlhMIzW"
      },
      "cell_type": "markdown",
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how `tf.distribute.Strategy` can be used for distributed multi-worker training with `tf.estimator`.  If you write your code using `tf.estimator`, and you're interested in scaling beyond a single machine with high performance, this tutorial is for you.\n",
        "\n",
        "Before getting started, please read the [`tf.distribute.Strategy` guide](../../guide/distribute_strategy.ipynb).  The [multi-GPU training tutorial](./keras.ipynb) is also relevant, because this tutorial uses the same model.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MUXex9ctTuDB"
      },
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "First, setup TensorFlow and the necessary imports."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "IqR2PQG4ZaZ0",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import, division, print_function"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bnYxvfLD-LW-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!pip install tensorflow-gpu==2.0.0-alpha0\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import os, json"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "hPBuZUNSZmrQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Input Function\n",
        "\n",
        "This tutorial uses the MNIST dataset from [TensorFlow Datasets](https://www.tensorflow.org/datasets).  The code here is similar to the [multi-GPU training tutorial](./keras.ipynb) with one key difference: when using Estimator for multi-worker training, it is necessary to shard the dataset by the number of workers to ensure model convergence.  The input data is sharded by worker index, so that each worker processes 1/`num_workers` distinct portions of the dataset."
      ]
    },
    {
      "metadata": {
        "id": "dma_wUAxZqo2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def input_fn(mode, input_context=None):\n",
        "  datasets, ds_info = tfds.load(name='mnist',\n",
        "                                with_info=True,\n",
        "                                as_supervised=True)\n",
        "  mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else\n",
        "                   datasets['test'])\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255\n",
        "    return image, label\n",
        "\n",
        "  if input_context:\n",
        "    mnist_dataset = mnist_dataset.shard(input_context.num_input_pipelines,\n",
        "                                        input_context.input_pipeline_id)\n",
        "  return mnist_dataset.map(scale).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
      ],
      "execution_count": 0,
{'fraction_non_alphanumeric': 0.1656142594516825, 'fraction_numerical': 0.0092718407801135, 'mean_word_length': 2.012761858897183, 'pattern_counts': {'":': 140, '<': 37, '<?xml version=': 0, '>': 37, 'https://': 16, 'lorem ipsum': 0, 'www.': 12, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '409407', 'n_tokens_mistral': 4178, 'n_tokens_neox': 4065, 'n_words': 879}