docs/tutorials/federated_learning_for_image_classification.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qN8P0AnTnAhh"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "form",
        "colab": {},
        "colab_type": "code",
        "id": "p8SrVqkmnDQv"
      },
      "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": {
        "colab_type": "text",
        "id": "AftvNA5VMemJ"
      },
      "source": [
        "# Federated Learning for Image Classification"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "coAumH42q9nz"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/federated_learning_for_image_classification\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/v0.14.0/docs/tutorials/federated_learning_for_image_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/v0.14.0/docs/tutorials/federated_learning_for_image_classification.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Zs2LgZBOMt4M"
      },
      "source": [
        "**NOTE**: This colab has been verified to work with the [latest released version](https://github.com/tensorflow/federated#compatibility) of the `tensorflow_federated` pip package, but the Tensorflow Federated project is still in pre-release development and may not work on `master`.\n",
        "\n",
        "In this tutorial, we use the classic MNIST training example to introduce the\n",
        "Federated Learning (FL) API layer of TFF, `tff.learning` - a set of\n",
        "higher-level interfaces that can be used to perform common types of federated\n",
        "learning tasks, such as federated training, against user-supplied models\n",
        "implemented in TensorFlow.\n",
        "\n",
        "This tutorial, and the Federated Learning API, are intended primarily for users\n",
        "who want to plug their own TensorFlow models into TFF, treating the latter\n",
        "mostly as a black box. For a more in-depth understanding of TFF and how to\n",
        "implement your own federated learning algorithms, see the tutorials on the FC Core API - [Custom Federated Algorithms Part 1](custom_federated_algorithms_1.ipynb) and [Part 2](custom_federated_algorithms_2.ipynb).\n",
        "\n",
        "For more on `tff.learning`, continue with the\n",
        "[Federated Learning for Text Generation](federated_learning_for_text_generation.ipynb),\n",
        "tutorial which in addition to covering recurrent models, also demonstrates loading a\n",
        "pre-trained serialized Keras model for refinement with federated learning\n",
        "combined with evaluation using Keras."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MnUwFbCAKB2r"
      },
      "source": [
        "## Before we start\n",
        "\n",
        "Before we start, please run the following to make sure that your environment is\n",
        "correctly setup. If you don't see a greeting, please refer to the\n",
        "[Installation](../install.md) guide for instructions. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZrGitA_KnRO0"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow_federated\n",
        "\n",
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 15063,
          "status": "ok",
          "timestamp": 1585108114376,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "8BKyHkMxKHfV",
        "outputId": "0fef9a2b-ed91-4393-a12c-25b1f7be3329"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "b'Hello, World!'"
            ]
          },
          "execution_count": 1,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import collections\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff\n",
        "\n",
        "tf.compat.v1.enable_v2_behavior()\n",
        "\n",
        "np.random.seed(0)\n",
        "\n",
        "tff.federated_computation(lambda: 'Hello, World!')()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5Cyy2AWbLMKj"
      },
      "source": [
        "## Preparing the input data\n",
        "\n",
        "Let's start with the data. Federated learning requires a federated data set,\n",
        "i.e., a collection of data from multiple users. Federated data is typically\n",
        "non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables),\n",
        "which poses a unique set of challenges.\n",
        "\n",
        "In order to facilitate experimentation, we seeded the TFF repository with a few\n",
        "datasets, including a federated version of MNIST that contains a version of the [original NIST dataset](https://www.nist.gov/srd/nist-special-database-19) that has been re-processed using [Leaf](https://github.com/TalwalkarLab/leaf) so that the data is keyed by the original writer of the digits. Since each writer has a unique style, this dataset exhibits the kind of non-i.i.d. behavior expected of federated datasets.\n",
        "\n",
        "Here's how we can load it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NayDhCX6SjwE"
      },
      "outputs": [],
      "source": [
        "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yeX8BKgPfeFw"
      },
      "source": [
        "The data sets returned by `load_data()` are instances of\n",
        "`tff.simulation.ClientData`, an interface that allows you to enumerate the set\n",
        "of users, to construct a `tf.data.Dataset` that represents the data of a\n",
        "particular user, and to query the structure of individual elements. Here's how\n",
        "you can use this interface to explore the content of the data set. Keep in mind\n",
        "that while this interface allows you to iterate over clients ids, this is only a\n",
        "feature of the simulation data. As you will see shortly, client identities are\n",
        "not used by the federated learning framework - their only purpose is to allow\n",
        "you to select subsets of the data for simulations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 456,
          "status": "ok",
          "timestamp": 1585108123118,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "kN4-U5nJgKig",
        "outputId": "043bb0c4-cd59-45bf-d861-ebc357af43fd"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "3383"
            ]
          },
          "execution_count": 3,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(emnist_train.client_ids)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 340,
          "status": "ok",
          "timestamp": 1585108124338,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "ZyCzIrSegT62",
        "outputId": "7fc4362f-e3c8-475e-deeb-8fdf2725b83f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "OrderedDict([('pixels', TensorSpec(shape=(28, 28), dtype=tf.float32, name=None)), ('label', TensorSpec(shape=(), dtype=tf.int32, name=None))])"
            ]
          },
          "execution_count": 4,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "emnist_train.element_type_structure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 581,
          "status": "ok",
          "timestamp": 1585108125831,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "EsvSXGEMgd9G",
        "outputId": "62702005-9bda-4e5f-bdaf-c61f0525f787"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "1"
            ]
          },
          "execution_count": 5,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "example_dataset = emnist_train.create_tf_dataset_for_client(\n",
        "    emnist_train.client_ids[0])\n",
        "\n",
        "example_element = next(iter(example_dataset))\n",
        "\n",
        "example_element['label'].numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 269
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 671,
          "status": "ok",
          "timestamp": 1585085538513,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "OmLV0nfMg98V",
        "outputId": "753e0d92-f882-4923-c35f-13ca4308b0ab"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBo\ndHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADbBJREFUeJzt3V9ok/cex/FP2uKdwsSmS61bVlq2\nrH8sNiIDKWgXNwqraGHoHFTqDMguhrqN7modDC0Mh4JemM2L4oXCLrQwZ1FkbqMwyuPsYCqjFMtq\nDW06C1MvptbfuTicnuPRJF2af/p9v0Bo8yQ+X574NmmfJD+fc84JgDklhR4AQGEQP2AU8QNGET9g\nFPEDRhE/YBTxA0YRP2AU8QNGleVzZ8uWLVMwGMznLgFTxsbGND09Pa/rLij+gYEBffDBB5qdndV7\n772n7u7ulNcPBoPyPG8huwSQQjgcnvd1M37aPzs7q/fff19nz57V1atXdeLECV29ejXTvw5AnmUc\n/9DQkGpqalRdXa1FixZpy5Yt6u/vz+ZsAHIo4/gnJia0YsWKue+rqqo0MTHx2PVisZjC4bDC4bAS\niUSmuwOQZRnH/6R3Avt8vscui0aj8jxPnuepvLw8090ByLKM46+qqtL4+Pjc9zdu3FBlZWVWhgKQ\nexnHv3r1ao2MjOj69eu6d++eTp48qfb29mzOBiCHMj7VV1ZWpsOHD+uNN97Q7Oysurq6VFdXl83Z\nAOTQgs7zt7W1qa2tLVuzAMgjXt4LGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8Y\nRfyAUcQPGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFHEDxhF\n/IBRxA8YtaBVeoPBoBYvXqzS0lKVlZXJ87xszQUgxxYUvyR9//33WrZsWTZmAZBHPO0HjFpQ/D6f\nTxs2bFBzc7NisVi2ZgKQBwt62j84OKjKykpNTU0pEonolVdeUUtLyyPXicVic/8xJBKJhewOQBYt\n6JG/srJSkuT3+7Vp0yYNDQ09dp1oNCrP8+R5nsrLyxeyOwBZlHH8d+/e1e3bt+e+PnfunOrr67M2\nGIDcyvhp/+TkpDZt2iRJevDggd555x29+eabWRsMQG5lHH91dbV+/fXXbM4CII841QcYRfyAUcQP\nGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8Y\nRfyAUcQPGEX8gFELXqUXz7aHDx+m3F5SwuPH04p7DjCK+AGjiB8wivgBo4gfMIr4AaOIHzAqbfxd\nXV3y+/2qr6+fu+zWrVuKRCKqra1VJBLRzMxMTodE7jjnUv4pKSlJ+QdPr7T33vbt2zUwMPDIZb29\nvWptbdXIyIhaW1vV29ubswEB5Eba+FtaWrR06dJHLuvv71dnZ6ckqbOzU6dPn87NdAByJqPnbZOT\nkwoEApKkQCCgqamprA4FIPdy/tr+WCymWCwmSUokErneHYB5yuiRv6KiQvF4XJIUj8fl9/uTXjca\njcrzPHmep/Ly8symBJB1GcXf3t6uvr4+SVJfX582btyY1aEA5F7a+Ldu3arXXntNv//+u6qqqnTs\n2DF1d3fr/Pnzqq2t1fnz59Xd3Z2PWQFkUdqf+U+cOPHEyy9cuJD1YZB9s7OzKbeXlpam3B6NRlNu\nX7duXcrtW7duTbptobNhYXiVBmAU8QNGET9gFPEDRhE/YBTxA0bx0d3PAOdc0m0+ny/lbaenp1Nu\nP3nyZMrtn3zyScrtqfCW4MLi6ANGET9gFPEDRhE/YBTxA0YRP2AU8QNGcZ7/GZDqrbFlZanv4i+/\n/DLl9ueeey7l9pdeeinl9lSz8ZbdwuKRHzCK+AGjiB8wivgBo4gfMIr4AaOIHzCK8/xPgYcPH6bc\nnup8+R9//JHytvv370+5/fz58ym3p/osARQ3HvkBo4gfMIr4AaOIHzCK+AGjiB8wivgBo9Ke5+/q\n6tK3334rv9+v3377TZLU09Ojr776SuXl5ZKkffv2qa2tLbeTGpbuXHqqz7/fvHlzytvu3r075fbX\nX3895XaW2X56pX3k3759uwYGBh67fPfu3RoeHtbw8DDhA0+htPG3tLRo6dKl+ZgFQB5l/DP/4cOH\n1djYqK6uLs3MzGRzJgB5kFH8u3bt0ujoqIaHhxUIBLR3796k143FYgqHwwqHw0okEhkPCiC7Moq/\noqJCpaWlKikp0c6dOzU0NJT0utFoVJ7nyfO8uV8QAii8jOKPx+NzX586dUr19fVZGwhAfqQ91bd1\n61ZdvHhR09PTqqqq0meffaaLFy9qeHhYPp9PwWBQR48ezcesALLI5/L4huxwOCzP8/K1O0i6du1a\nyu2hUCjl9nSfJZDqNQbIv3/SGPccYBTxA0YRP2AU8QNGET9gFPEDRvHR3c+4dKfyFvJ2YTzduGcB\no4gfMIr4AaOIHzCK+AGjiB8wivgBozjP/4zjLblIhnseMIr4AaOIHzCK+AGjiB8wivgBo4gfMIrz\n/M84zuMjGf5lAEYRP2AU8QNGET9gFPEDRhE/YBTxA0aljX98fFzr1q1TKBRSXV2dDh06JEm6deuW\nIpGIamtrFYlENDMzk/NhAWRP2vjLysp04MABXbt2TT///LOOHDmiq1evqre3V62trRoZGVFra6t6\ne3vzMS+ALEkbfyAQ0KpVqyRJixcvVigU0sTEhPr7+9XZ2SlJ6uzs1OnTp3M7KYCs+kc/84+Njeny\n5ctas2aNJicnFQgEJP37P4ipqamcDAggN+b92v47d+6oo6NDBw8e1JIlS+a9g1gsplgsJklKJBL/\nfEIAOTGvR/779++ro6ND27Zt0+bNmyVJFRUVisfjkqR4PC6/3//E20ajUXmeJ8/zVF5enqWxASxU\n2vidc9qxY4dCoZD27Nkzd3l7e7v6+vokSX19fdq4cWPupgSQdWmf9g8ODur48eNqaGhQU1OTJGnf\nvn3q7u7W22+/rWPHjumFF17QN998k/NhAWRP2vjXrl2bdA33CxcuZH0gAPnBK/wAo4gfMIr4AaOI\nHzCK+AGjiB8wivgBo4gfMIr4AaOIHzCK+AGjiB8wivgBo4gfMIr4AaOIHzCK+AGjiB8wivgBo4gf\nMIr4AaOIHzCK+AGjiB8wivgBo4gfMIr4AaOIHzCK+AGjiB8wKm384+PjWrdunUKhkOrq6nTo0CFJ\nUk9Pj5YvX66mpiY1NTXpu+++y/mwALKnLO0Vysp04MABrVq1Srdv31Zzc7MikYgkaffu3frwww9z\nPiSA7EsbfyAQUCAQkCQtXrxYoVBIExMTOR8MQG79o5/5x8bGdPnyZa1Zs0aSdPjwYTU2Nqqrq0sz\nMzNPvE0sFlM4HFY4HFYikVj4xACyYt7x37lzRx0dHTp48KCWLFmiXbt2aXR0VMPDwwoEAtq7d+8T\nbxeNRuV5njzPU3l5edYGB7Aw84r//v376ujo0LZt27R582ZJUkVFhUpLS1VSUqKdO3dqaGgop4MC\nyK608TvntGPHDoVCIe3Zs2fu8ng8Pvf1qVOnVF9fn5sJAeRE2l/4DQ4O6vjx42poaFBTU5Mkad++\nfTpx4oSGh4fl8/kUDAZ19OjRnA8LIHvSxr927Vo55x67vK2tLScDAcgPXuEHGEX8gFHEDxhF/IBR\nxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8YRfyAUcQPGEX8gFE+96T36+bIsmXLFAwG575PJBJF+9Fe\nxTpbsc4lMVumsjnb2NiYpqen53XdvMb//8LhsDzPK9TuUyrW2Yp1LonZMlWo2XjaDxhF/IBRpT09\nPT2FHKC5ubmQu0+pWGcr1rkkZstUIWYr6M/8AAqHp/2AUQWJf2BgQC+//LJqamrU29tbiBGSCgaD\ncx9THg6HCzpLV1eX/H7/I2si3Lp1S5FIRLW1tYpEIkmXSSvEbMWycnOylaULfeyKbsVrl2cPHjxw\n1dXVbnR01P3999+usbHRXblyJd9jJPXiiy+6RCJR6DGcc8798MMP7tKlS66urm7uso8++sjt37/f\nOefc/v373ccff1w0s3366afuiy++KMg8/+vmzZvu0qVLzjnn/vrrL1dbW+uuXLlS8GOXbK5CHbe8\nP/IPDQ2ppqZG1dXVWrRokbZs2aL+/v58j/FUaGlp0dKlSx+5rL+/X52dnZKkzs5OnT59uhCjPXG2\nYhEIBLRq1SpJj64sXehjl2yuQsl7/BMTE1qxYsXc91VVVUW15LfP59OGDRvU3NysWCxW6HEeMzk5\nObdkeiAQ0NTUVIEnetR8Vm7Op/9dWbqYjl0mK15nW97jd084ueDz+fI9RlKDg4P65ZdfdPbsWR05\nckQ//vhjoUd6asx35eZ8+f+VpYtFpiteZ1ve46+qqtL4+Pjc9zdu3FBlZWW+x0jqP7P4/X5t2rSp\n6FYfrqiomFskNR6Py+/3F3ii/yqmlZuTrSxd6GNXTCte5z3+1atXa2RkRNevX9e9e/d08uRJtbe3\n53uMJ7p7965u37499/W5c+eKbvXh9vZ29fX1SZL6+vq0cePGAk/0X8WycrNLsrJ0oY9dsrkKdtzy\n/itG59yZM2dcbW2tq66udp9//nkhRnii0dFR19jY6BobG92rr75a8Nm2bNninn/+eVdWVuaWL1/u\nvv76azc9Pe3Wr1/vampq3Pr1692ff/5ZNLO9++67rr6+3jU0NLi33nrL3bx5syCz/fTTT06Sa2ho\ncCtXrnQrV650Z86cKfixSzZXoY4br/ADjOIVfoBRxA8YRfyAUcQPGEX8gFHEDxhF/IBRxA8Y9S85\nOQzW0CUnRwAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "\u003cFigure size 600x400 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "plt.imshow(example_element['pixels'].numpy(), cmap='gray', aspect='equal')\n",
        "plt.grid(False)\n",
        "_ = plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "lMd01egqy9we"
      },
      "source": [
        "Since the data is already a `tf.data.Dataset`,  preprocessing can be accomplished using Dataset transformations. Here, we flatten the `28x28` images\n",
        "into `784`-element arrays, shuffle the individual examples, organize them into batches, and renames the features\n",
        "from `pixels` and `label` to `x` and `y` for use with Keras. We also throw in a\n",
        "`repeat` over the data set to run several epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "cyG_BMraSuu_"
      },
      "outputs": [],
      "source": [
        "NUM_CLIENTS = 10\n",
        "NUM_EPOCHS = 5\n",
        "BATCH_SIZE = 20\n",
        "SHUFFLE_BUFFER = 100\n",
        "PREFETCH_BUFFER=10\n",
        "\n",
        "def preprocess(dataset):\n",
        "\n",
        "  def batch_format_fn(element):\n",
        "    \"\"\"Flatten a batch `pixels` and return the features as an `OrderedDict`.\"\"\"\n",
        "    return collections.OrderedDict(\n",
        "        x=tf.reshape(element['pixels'], [-1, 784]),\n",
        "        y=tf.reshape(element['label'], [-1, 1]))\n",
        "\n",
        "  return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(\n",
        "      BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m9LXykN_jlJw"
      },
      "source": [
        "Let's verify this worked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 504
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 620,
          "status": "ok",
          "timestamp": 1585085541675,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "VChB7LMQjkYz",
        "outputId": "8298e227-ae95-4124-f8f9-1bbc72ac4ea3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "OrderedDict([('x', array([[1., 1., 1., ..., 1., 1., 1.],\n",
              "       [1., 1., 1., ..., 1., 1., 1.],\n",
              "       [1., 1., 1., ..., 1., 1., 1.],\n",
              "       ...,\n",
              "       [1., 1., 1., ..., 1., 1., 1.],\n",
              "       [1., 1., 1., ..., 1., 1., 1.],\n",
              "       [1., 1., 1., ..., 1., 1., 1.]], dtype=float32)), ('y', array([[6],\n",
              "       [6],\n",
              "       [7],\n",
              "       [9],\n",
              "       [1],\n",
              "       [2],\n",
              "       [4],\n",
              "       [0],\n",
              "       [1],\n",
              "       [4],\n",
              "       [7],\n",
              "       [8],\n",
              "       [6],\n",
              "       [0],\n",
              "       [6],\n",
              "       [7],\n",
              "       [5],\n",
              "       [3],\n",
              "       [2],\n",
              "       [7]], dtype=int32))])"
            ]
          },
          "execution_count": 8,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "preprocessed_example_dataset = preprocess(example_dataset)\n",
        "\n",
        "sample_batch = tf.nest.map_structure(lambda x: x.numpy(),\n",
        "                                     next(iter(preprocessed_example_dataset)))\n",
        "\n",
        "sample_batch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JGsMvRQt9Agl"
      },
      "source": [
        "We have almost all the building blocks in place to construct federated data\n",
        "sets.\n",
        "\n",
        "One of the ways to feed federated data to TFF in a simulation is simply as a\n",
        "Python list, with each element of the list holding the data of an individual\n",
        "user, whether as a list or as a `tf.data.Dataset`. Since we already have\n",
        "an interface that provides the latter, let's use it.\n",
        "\n",
        "Here's a simple helper function that will construct a list of datasets from the\n",
        "given set of users as an input to a round of training or evaluation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "_PHMvHAI9xVc"
      },
      "outputs": [],
      "source": [
        "def make_federated_data(client_data, client_ids):\n",
        "  return [\n",
        "      preprocess(client_data.create_tf_dataset_for_client(x))\n",
        "      for x in client_ids\n",
        "  ]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0M9PfjOtAVqw"
      },
      "source": [
        "Now, how do we choose clients?\n",
        "\n",
        "In a typical federated training scenario, we are dealing with potentially a very\n",
        "large population of user devices, only a fraction of which may be available for\n",
        "training at a given point in time. This is the case, for example, when the\n",
        "client devices are mobile phones that participate in training only when plugged\n",
        "into a power source, off a metered network, and otherwise idle.\n",
        "\n",
        "Of course, we are in a simulation environment, and all the data is locally\n",
        "available. Typically then, when running simulations, we would simply sample a\n",
        "random subset of the clients to be involved in each round of training, generally\n",
        "different in each round.\n",
        "\n",
        "That said, as you can find out by studying the paper on the\n",
        "[Federated Averaging](https://arxiv.org/abs/1602.05629) algorithm, achieving convergence in a system with randomly sampled\n",
        "subsets of clients in each round can take a while, and it would be impractical\n",
        "to have to run hundreds of rounds in this interactive tutorial.\n",
        "\n",
        "What we'll do instead is sample the set of clients once, and\n",
        "reuse the same set across rounds to speed up convergence (intentionally\n",
        "over-fitting to these few user's data). We leave it as an exercise for the\n",
        "reader to modify this tutorial to simulate random sampling - it is fairly easy to\n",
        "do (once you do, keep in mind that getting the model to converge may take a\n",
        "while)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 54
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 599,
          "status": "ok",
          "timestamp": 1585085547013,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "GZ6NYHxB8xer",
        "outputId": "e699025e-5505-4d1a-fa9e-393e6f517476"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of client datasets: 10\n",
            "First dataset: \u003cDatasetV1Adapter shapes: OrderedDict([(x, (None, 784)), (y, (None, 1))]), types: OrderedDict([(x, tf.float32), (y, tf.int32)])\u003e\n"
          ]
        }
      ],
      "source": [
        "sample_clients = emnist_train.client_ids[0:NUM_CLIENTS]\n",
        "\n",
        "federated_train_data = make_federated_data(emnist_train, sample_clients)\n",
        "\n",
        "print('Number of client datasets: {l}'.format(l=len(federated_train_data)))\n",
        "print('First dataset: {d}'.format(d=federated_train_data[0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "HOxq4tbi9m8-"
      },
      "source": [
        "## Creating a model with Keras\n",
        "\n",
        "If you are using Keras, you likely already have code that constructs a Keras\n",
        "model. Here's an example of a simple model that will suffice for our needs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "LYCsJGJFWbqt"
      },
      "outputs": [],
      "source": [
        "def create_keras_model():\n",
        "  return tf.keras.models.Sequential([\n",
        "      tf.keras.layers.Input(shape=(784,)),\n",
        "      tf.keras.layers.Dense(10, kernel_initializer='zeros'),\n",
        "      tf.keras.layers.Softmax(),\n",
        "  ])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NHdraKFH4OU2"
      },
      "source": [
        "**Note:** we do not compile the model yet. The loss, metrics, and optimizers are introduced later.\n",
        "\n",
        "In order to use any model with TFF, it needs to be wrapped in an instance of the\n",
        "`tff.learning.Model` interface, which exposes methods to stamp the model's\n",
        "forward pass, metadata properties, etc., similarly to Keras, but also introduces\n",
        "additional elements, such as ways to control the process of computing federated\n",
        "metrics. Let's not worry about this for now; if you have a Keras model like the\n",
        "one we've just defined above, you can have TFF wrap it for you by invoking\n",
        "`tff.learning.from_keras_model`, passing the model and a sample data batch as\n",
        "arguments, as shown below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Q3ynrxd53HzY"
      },
      "outputs": [],
      "source": [
        "def model_fn():\n",
        "  # We _must_ create a new model here, and _not_ capture it from an external\n",
        "  # scope. TFF will call this within different graph contexts.\n",
        "  keras_model = create_keras_model()\n",
        "  return tff.learning.from_keras_model(\n",
        "      keras_model,\n",
        "      input_spec=preprocessed_example_dataset.element_spec,\n",
        "      loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
        "      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XJ5E3O18_JZ6"
      },
      "source": [
        "## Training the model on federated data\n",
        "\n",
        "Now that we have a model wrapped as `tff.learning.Model` for use with TFF, we\n",
        "can let TFF construct a Federated Averaging algorithm by invoking the helper\n",
        "function `tff.learning.build_federated_averaging_process`, as follows.\n",
        "\n",
        "Keep in mind that the argument needs to be a constructor (such as `model_fn`\n",
        "above), not an already-constructed instance, so that the construction of your\n",
        "model can happen in a context controlled by TFF (if you're curious about the\n",
        "reasons for this, we encourage you to read the follow-up tutorial on\n",
        "[custom algorithms](custom_federated_algorithms_1.ipynb)).\n",
        "\n",
        "One critical note on the Federated Averaging algorithm below, there are **2**\n",
        "optimizers: a _client_optimizer_ and a _server_optimizer_. The\n",
        "_client_optimizer_ is only used to compute local model updates on each client.\n",
        "The _server_optimizer_ applies the averaged update to the global model at the\n",
        "server. In particular, this means that the choice of optimizer and learning rate\n",
        "used may need to be different than the ones you have used to train the model on\n",
        "a standard i.i.d. dataset. We recommend starting with regular SGD, possibly with\n",
        "a smaller learning rate than usual. The learning rate we use has not been\n",
        "carefully tuned, feel free to experiment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "sk6mjOfycX5N"
      },
      "outputs": [],
      "source": [
        "iterative_process = tff.learning.build_federated_averaging_process(\n",
        "    model_fn,\n",
        "    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),\n",
        "    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "f8FpvN2n67sm"
      },
      "source": [
        "What just happened? TFF has constructed a pair of *federated computations* and\n",
        "packaged them into a `tff.templates.IterativeProcess` in which these computations\n",
        "are available as a pair of properties `initialize` and `next`.\n",
        "\n",
        "In a nutshell, *federated computations* are programs in TFF's internal language\n",
        "that can express various federated algorithms (you can find more about this in\n",
        "the [custom algorithms](custom_federated_algorithms_1.ipynb) tutorial). In this\n",
        "case, the two computations generated and packed into `iterative_process`\n",
        "implement [Federated Averaging](https://arxiv.org/abs/1602.05629).\n",
        "\n",
        "It is a goal of TFF to define computations in a way that they could be executed\n",
        "in real federated learning settings, but currently only local execution\n",
        "simulation runtime is implemented. To execute a computation in a simulator, you\n",
        "simply invoke it like a Python function. This default interpreted environment is\n",
        "not designed for high performance, but it will suffice for this tutorial; we\n",
        "expect to provide higher-performance simulation runtimes to facilitate\n",
        "larger-scale research in future releases.\n",
        "\n",
        "Let's start with the `initialize` computation. As is the case for all federated\n",
        "computations, you can think of it as a function. The computation takes no\n",
        "arguments, and returns one result - the representation of the state of the\n",
        "Federated Averaging process on the server. While we don't want to dive into the\n",
        "details of TFF, it may be instructive to see what this state looks like. You can\n",
        "visualize it as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 400,
          "status": "ok",
          "timestamp": 1585085572354,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "Z4pcfWsUBp_5",
        "outputId": "acb72fc0-ce78-41b7-8f44-b89a26c6c8a4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'( -\u003e \u003cmodel=\u003ctrainable=\u003cfloat32[784,10],float32[10]\u003e,non_trainable=\u003c\u003e\u003e,optimizer_state=\u003cint64\u003e,delta_aggregate_state=\u003c\u003e,model_broadcast_state=\u003c\u003e\u003e@SERVER)'"
            ]
          },
          "execution_count": 15,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(iterative_process.initialize.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "v1gbHQ_7BiyT"
      },
      "source": [
        "While the above type signature may at first seem a bit cryptic, you can\n",
        "recognize that the server state consists of a `model` (the initial model\n",
        "parameters for MNIST that will be distributed to all devices), and\n",
        "`optimizer_state` (additional information maintained by the server, such as the\n",
        "number of rounds to use for hyperparameter schedules, etc.).\n",
        "\n",
        "Let's invoke the `initialize` computation to construct the server state."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6cagCWlZmcch"
      },
      "outputs": [],
      "source": [
        "state = iterative_process.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TjjxTx9e_rMd"
      },
      "source": [
        "The second of the pair of federated computations, `next`, represents a single\n",
        "round of Federated Averaging, which consists of pushing the server state\n",
        "(including the model parameters) to the clients, on-device training on their\n",
        "local data, collecting and averaging model updates, and producing a new updated\n",
        "model at the server.\n",
        "\n",
        "Conceptually, you can think of `next` as having a functional type signature that\n",
        "looks as follows.\n",
        "\n",
        "```\n",
        "SERVER_STATE, FEDERATED_DATA -\u003e SERVER_STATE, TRAINING_METRICS\n",
        "```\n",
        "\n",
        "In particular, one should think about `next()` not as being a function that runs on a server, but rather being a declarative functional representation of the entire decentralized computation - some of the inputs are provided by the server (`SERVER_STATE`), but each participating device contributes its own local dataset.\n",
        "\n",
        "Let's run a single round of training and visualize the results. We can use the\n",
        "federated data we've already generated above for a sample of users."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 36
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 3949,
          "status": "ok",
          "timestamp": 1585085579305,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "F3M_W9dDE6Tm",
        "outputId": "fc87fca5-c23c-4040-a534-08c0e55b042e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round  1, metrics=\u003csparse_categorical_accuracy=0.1281892955303192,loss=3.112910270690918,keras_training_time_client_sum_sec=0.0\u003e\n"
          ]
        }
      ],
      "source": [
        "state, metrics = iterative_process.next(state, federated_train_data)\n",
        "print('round  1, metrics={}'.format(metrics))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UmhReXt9G4A5"
      },
      "source": [
        "Let's run a few more rounds. As noted earlier, typically at this point you would\n",
        "pick a subset of your simulation data from a new randomly selected sample of\n",
        "users for each round in order to simulate a realistic deployment in which users\n",
        "continuously come and go, but in this interactive notebook, for the sake of\n",
        "demonstration we'll just reuse the same users, so that the system converges\n",
        "quickly."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 186
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 21922,
          "status": "ok",
          "timestamp": 1585085602179,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "qrJkQuCRJP9C",
        "outputId": "9d95386f-9c67-488b-fb4d-c4b638b475e4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round  2, metrics=\u003csparse_categorical_accuracy=0.1366255134344101,loss=3.0699315071105957,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  3, metrics=\u003csparse_categorical_accuracy=0.14938271045684814,loss=2.967233657836914,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  4, metrics=\u003csparse_categorical_accuracy=0.17860081791877747,loss=2.7275609970092773,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  5, metrics=\u003csparse_categorical_accuracy=0.20637859404087067,loss=2.601724863052368,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  6, metrics=\u003csparse_categorical_accuracy=0.2100823074579239,loss=2.5941531658172607,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  7, metrics=\u003csparse_categorical_accuracy=0.244650200009346,loss=2.3704617023468018,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  8, metrics=\u003csparse_categorical_accuracy=0.2720164656639099,loss=2.225743293762207,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round  9, metrics=\u003csparse_categorical_accuracy=0.28641974925994873,loss=2.2297680377960205,keras_training_time_client_sum_sec=0.0\u003e\n",
            "round 10, metrics=\u003csparse_categorical_accuracy=0.31975308060646057,loss=2.0716400146484375,keras_training_time_client_sum_sec=0.0\u003e\n"
          ]
        }
      ],
      "source": [
        "NUM_ROUNDS = 11\n",
        "for round_num in range(2, NUM_ROUNDS):\n",
        "  state, metrics = iterative_process.next(state, federated_train_data)\n",
        "  print('round {:2d}, metrics={}'.format(round_num, metrics))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "joHYzn9jcs0Y"
      },
      "source": [
        "Training loss is decreasing after each round of federated training, indicating\n",
        "the model is converging. There are some important caveats with these training\n",
        "metrics, however, see the section on *Evaluation* later in this tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ruSHJl1IjhNf"
      },
      "source": [
        "##Displaying model metrics in TensorBoard\n",
        "Next, let's visualize the metrics from these federated computations using Tensorboard.\n",
        "\n",
        "Let's start by creating the directory and the corresponding summary writer to write the metrics to.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "E3QUBK41lWDW"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "logdir = \"/tmp/logs/scalars/training/\"\n",
        "summary_writer = tf.summary.create_file_writer(logdir)\n",
        "state = iterative_process.initialize()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "w-2aGxUlzS_J"
      },
      "source": [
        "Plot the relevant scalar metrics with the same summary writer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "JZtr4_8lzN-V"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "with summary_writer.as_default():\n",
        "  for round_num in range(1, NUM_ROUNDS):\n",
        "    state, metrics = iterative_process.next(state, federated_train_data)\n",
        "    for name, value in metrics._asdict().items():\n",
        "      tf.summary.scalar(name, value, step=round_num)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "iUouyAHG0Mk8"
      },
      "source": [
        "Start TensorBoard with the root log directory specified above. It can take a few seconds for the data to load."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "urYYcmA9089p"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "%tensorboard --logdir /tmp/logs/scalars/ --port=0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZMcV15W7b1wG"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "# Run this this cell to clean your directory of old output for future graphs from this directory.\n",
        "!rm -R /tmp/logs/scalars/*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jejrFEVP1EDs"
      },
      "source": [
        "In order to view evaluation metrics the same way, you can create a separate eval folder, like \"logs/scalars/eval\", to write to TensorBoard."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "T4hneAcb-F2l"
      },
      "source": [
        "## Customizing the model implementation\n",
        "\n",
        "Keras is the [recommended high-level model API for TensorFlow](https://medium.com/tensorflow/standardizing-on-keras-guidance-on-high-level-apis-in-tensorflow-2-0-bad2b04c819a), and we encourage using Keras models (via \n",
        "`tff.learning.from_keras_model` or\n",
        "`tff.learning.from_compiled_keras_model`) in TFF whenever possible.\n",
        "\n",
        "However, `tff.learning` provides a lower-level model interface, `tff.learning.Model`, that exposes the minimal functionality necessary for using a model for federated learning. Directly implementing this interface (possibly still using building blocks like `tf.keras.layers`) allows for maximum customization without modifying the internals of the federated learning algorithms.\n",
        "\n",
        "So let's do it all over again from scratch.\n",
        "\n",
        "### Defining model variables, forward pass, and metrics\n",
        "\n",
        "The first step is to identify the TensorFlow variables we're going to work with.\n",
        "In order to make the following code more legible, let's define a data structure\n",
        "to represent the entire set. This will include variables such as `weights` and\n",
        "`bias` that we will train, as well as variables that will hold various\n",
        "cumulative statistics and counters we will update during training, such as\n",
        "`loss_sum`, `accuracy_sum`, and `num_examples`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "uqRD72WQC4u1"
      },
      "outputs": [],
      "source": [
        "MnistVariables = collections.namedtuple(\n",
        "    'MnistVariables', 'weights bias num_examples loss_sum accuracy_sum')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "nkJfDcY5oXii"
      },
      "source": [
        "Here's a method that creates the variables. For the sake of simplicity, we\n",
        "represent all statistics as `tf.float32`, as that will eliminate the need for\n",
        "type conversions at a later stage. Wrapping variable initializers as lambdas is\n",
        "a requirement imposed by\n",
        "[resource variables](https://www.tensorflow.org/api_docs/python/tf/enable_resource_variables)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "H3GQHLNqCfMU"
      },
      "outputs": [],
      "source": [
        "def create_mnist_variables():\n",
        "  return MnistVariables(\n",
        "      weights=tf.Variable(\n",
        "          lambda: tf.zeros(dtype=tf.float32, shape=(784, 10)),\n",
        "          name='weights',\n",
        "          trainable=True),\n",
        "      bias=tf.Variable(\n",
        "          lambda: tf.zeros(dtype=tf.float32, shape=(10)),\n",
        "          name='bias',\n",
        "          trainable=True),\n",
        "      num_examples=tf.Variable(0.0, name='num_examples', trainable=False),\n",
        "      loss_sum=tf.Variable(0.0, name='loss_sum', trainable=False),\n",
        "      accuracy_sum=tf.Variable(0.0, name='accuracy_sum', trainable=False))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SrdnR0fAre-Q"
      },
      "source": [
        "With the variables for model parameters and cumulative statistics in place, we\n",
        "can now define the forward pass method that computes loss, emits predictions,\n",
        "and updates the cumulative statistics for a single batch of input data, as\n",
        "follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZYSRAl-KCvC7"
      },
      "outputs": [],
      "source": [
        "def mnist_forward_pass(variables, batch):\n",
        "  y = tf.nn.softmax(tf.matmul(batch['x'], variables.weights) + variables.bias)\n",
        "  predictions = tf.cast(tf.argmax(y, 1), tf.int32)\n",
        "\n",
        "  flat_labels = tf.reshape(batch['y'], [-1])\n",
        "  loss = -tf.reduce_mean(\n",
        "      tf.reduce_sum(tf.one_hot(flat_labels, 10) * tf.math.log(y), axis=[1]))\n",
        "  accuracy = tf.reduce_mean(\n",
        "      tf.cast(tf.equal(predictions, flat_labels), tf.float32))\n",
        "\n",
        "  num_examples = tf.cast(tf.size(batch['y']), tf.float32)\n",
        "\n",
        "  variables.num_examples.assign_add(num_examples)\n",
        "  variables.loss_sum.assign_add(loss * num_examples)\n",
        "  variables.accuracy_sum.assign_add(accuracy * num_examples)\n",
        "\n",
        "  return loss, predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "-gm-yx2Mr_bl"
      },
      "source": [
        "Next, we define a function that returns a set of local metrics, again using TensorFlow. These are the values (in addition to model updates, which are handled automatically) that are eligible to be aggregated to the server in a federated learning or evaluation process.\n",
        "\n",
        "Here, we simply return the average `loss` and `accuracy`, as well as the\n",
        "`num_examples`, which we'll need to correctly weight the contributions from\n",
        "different users when computing federated aggregates."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "RkAZXhjGEekp"
      },
      "outputs": [],
      "source": [
        "def get_local_mnist_metrics(variables):\n",
        "  return collections.OrderedDict(\n",
        "      num_examples=variables.num_examples,\n",
        "      loss=variables.loss_sum / variables.num_examples,\n",
        "      accuracy=variables.accuracy_sum / variables.num_examples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9ywGs1G-s1o3"
      },
      "source": [
        "Finally, we need to determine how to aggregate the local metrics emitted by each\n",
        "device via `get_local_mnist_metrics`. This is the only part of the code that isn't written in TensorFlow  - it's a *federated computation* expressed in TFF. If you'd like to\n",
        "dig deeper, skim over the [custom algorithms](custom_federated_algorithms_1.ipynb)\n",
        "tutorial, but in most applications, you won't really need to; variants of the\n",
        "pattern shown below should suffice. Here's what it looks like:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "BMr2PwkfExFI"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation\n",
        "def aggregate_mnist_metrics_across_clients(metrics):\n",
        "  return collections.OrderedDict(\n",
        "      num_examples=tff.federated_sum(metrics.num_examples),\n",
        "      loss=tff.federated_mean(metrics.loss, metrics.num_examples),\n",
        "      accuracy=tff.federated_mean(metrics.accuracy, metrics.num_examples))\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2rXZ3Hg44aeN"
      },
      "source": [
        "The input `metrics` argument corresponds to the `OrderedDict` returned by `get_local_mnist_metrics` above, but critically the values are no longer `tf.Tensors` - they are \"boxed\" as `tff.Value`s, to make it clear you can no longer manipulate them using TensorFlow, but only using TFF's federated operators like `tff.federated_mean` and `tff.federated_sum`.  The returned\n",
        "dictionary of global aggregates defines the set of metrics which will be available on the server.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7MXGAuQRvmcp"
      },
      "source": [
        "### Constructing an instance of `tff.learning.Model`\n",
        "\n",
        "With all of the above in place, we are ready to construct a model representation\n",
        "for use with TFF similar to one that's generated for you when you let TFF ingest\n",
        "a Keras model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "blQGiTQFS9_r"
      },
      "outputs": [],
      "source": [
        "class MnistModel(tff.learning.Model):\n",
        "\n",
        "  def __init__(self):\n",
        "    self._variables = create_mnist_variables()\n",
        "\n",
        "  @property\n",
        "  def trainable_variables(self):\n",
        "    return [self._variables.weights, self._variables.bias]\n",
        "\n",
        "  @property\n",
        "  def non_trainable_variables(self):\n",
        "    return []\n",
        "\n",
        "  @property\n",
        "  def local_variables(self):\n",
        "    return [\n",
        "        self._variables.num_examples, self._variables.loss_sum,\n",
        "        self._variables.accuracy_sum\n",
        "    ]\n",
        "\n",
        "  @property\n",
        "  def input_spec(self):\n",
        "    return collections.OrderedDict(\n",
        "        x=tf.TensorSpec([None, 784], tf.float32),\n",
        "        y=tf.TensorSpec([None, 1], tf.int32))\n",
        "\n",
        "  @tf.function\n",
        "  def forward_pass(self, batch, training=True):\n",
        "    del training\n",
        "    loss, predictions = mnist_forward_pass(self._variables, batch)\n",
        "    num_exmaples = tf.shape(batch['x'])[0]\n",
        "    return tff.learning.BatchOutput(\n",
        "        loss=loss, predictions=predictions, num_examples=num_exmaples)\n",
        "\n",
        "  @tf.function\n",
        "  def report_local_outputs(self):\n",
        "    return get_local_mnist_metrics(self._variables)\n",
        "\n",
        "  @property\n",
        "  def federated_output_computation(self):\n",
        "    return aggregate_mnist_metrics_across_clients"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sMN1AszMwLHL"
      },
      "source": [
        "As you can see, the abstract methods and properties defined by\n",
        "`tff.learning.Model` corresponds to the code snippets in the preceding section\n",
        "that introduced the variables and defined the loss and statistics.\n",
        "\n",
        "Here are a few points worth highlighting:\n",
        "\n",
        "*   All state that your model will use must be captured as TensorFlow variables,\n",
        "    as TFF does not use Python at runtime (remember your code should be written\n",
        "    such that it can be deployed to mobile devices; see the\n",
        "    [custom algorithms](custom_federated_algorithms_1.ipynb) tutorial for a more\n",
        "    in-depth commentary on the reasons).\n",
        "*   Your model should describe what form of data it accepts (`input_spec`), as\n",
        "    in general, TFF is a strongly-typed environment and wants to determine type\n",
        "    signatures for all components. Declaring the format of your model's input is\n",
        "    an essential part of it.\n",
        "*   Although technically not required, we recommend wrapping all TensorFlow\n",
        "    logic (forward pass, metric calculations, etc.) as `tf.function`s,\n",
        "    as this helps ensure the TensorFlow can be serialized, and removes the need\n",
        "    for explicit control dependencies.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9DVhXk2Bu-GU"
      },
      "source": [
        "The above is sufficient for evaluation and algorithms like Federated SGD.\n",
        "However, for Federated Averaging, we need to specify how the model should train\n",
        "locally on each batch. We will specify a local optimizer when building the Federated Averaging algorithm."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hVBugKP3yw03"
      },
      "source": [
        "### Simulating federated training with the new model\n",
        "\n",
        "With all the above in place, the remainder of the process looks like what we've\n",
        "seen already - just replace the model constructor with the constructor of our\n",
        "new model class, and use the two federated computations in the iterative process\n",
        "you created to cycle through training rounds."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "FK3c8_leS9_t"
      },
      "outputs": [],
      "source": [
        "iterative_process = tff.learning.build_federated_averaging_process(\n",
        "    MnistModel,\n",
        "    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Jv_LiggwS9_u"
      },
      "outputs": [],
      "source": [
        "state = iterative_process.initialize()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 3366,
          "status": "ok",
          "timestamp": 1584130574170,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "PtOLElmzDPxs",
        "outputId": "971fbb55-c2e2-4b16-c312-78135c16cbef"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round  1, metrics=\u003cnum_examples=4860.0,loss=3.1408212184906006,accuracy=0.11687242984771729\u003e\n"
          ]
        }
      ],
      "source": [
        "state, metrics = iterative_process.next(state, federated_train_data)\n",
        "print('round  1, metrics={}'.format(metrics))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 170
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 20712,
          "status": "ok",
          "timestamp": 1584130594889,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "gFkv0yJEGhue",
        "outputId": "854047ec-e303-446a-f38f-8884bc904f42"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round  2, metrics=\u003cnum_examples=4860.0,loss=3.0044455528259277,accuracy=0.1366255134344101\u003e\n",
            "round  3, metrics=\u003cnum_examples=4860.0,loss=2.8893046379089355,accuracy=0.14753086864948273\u003e\n",
            "round  4, metrics=\u003cnum_examples=4860.0,loss=2.742488384246826,accuracy=0.1720164567232132\u003e\n",
            "round  5, metrics=\u003cnum_examples=4860.0,loss=2.52746844291687,accuracy=0.20843622088432312\u003e\n",
            "round  6, metrics=\u003cnum_examples=4860.0,loss=2.455655574798584,accuracy=0.23127572238445282\u003e\n",
            "round  7, metrics=\u003cnum_examples=4860.0,loss=2.369921922683716,accuracy=0.2522633671760559\u003e\n",
            "round  8, metrics=\u003cnum_examples=4860.0,loss=2.288294553756714,accuracy=0.27345678210258484\u003e\n",
            "round  9, metrics=\u003cnum_examples=4860.0,loss=2.18411922454834,accuracy=0.2890946567058563\u003e\n",
            "round 10, metrics=\u003cnum_examples=4860.0,loss=2.048980951309204,accuracy=0.32510289549827576\u003e\n"
          ]
        }
      ],
      "source": [
        "for round_num in range(2, 11):\n",
        "  state, metrics = iterative_process.next(state, federated_train_data)\n",
        "  print('round {:2d}, metrics={}'.format(round_num, metrics))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Iswqa2Uj7phq"
      },
      "source": [
        "To see these metrics within TensorBoard, refer to the steps listed above in \"Displaying model metrics in TensorBoard\"."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "m7lz59lMJ0kj"
      },
      "source": [
        "## Evaluation\n",
        "\n",
        "All of our experiments so far presented only federated training metrics - the\n",
        "average metrics over all batches of data trained across all clients in the\n",
        "round. This introduces the normal concerns about overfitting, especially since\n",
        "we used the same set of clients on each round for simplicity, but there is an\n",
        "additional notion of overfitting in training metrics specific to the Federated\n",
        "Averaging algorithm. This is easiest to see if we imagine each client had a\n",
        "single batch of data, and we train on that batch for many iterations (epochs).\n",
        "In this case, the local model will quickly exactly fit to that one batch, and so\n",
        "the local accuracy metric we average will approach 1.0. Thus, these training\n",
        "metrics can be taken as a sign that training is progressing, but not much more.\n",
        "\n",
        "To perform evaluation on federated data, you can construct another *federated\n",
        "computation* designed for just this purpose, using the\n",
        "`tff.learning.build_federated_evaluation` function, and passing in your model\n",
        "constructor as an argument. Note that unlike with Federated Averaging, where\n",
        "we've used `MnistTrainableModel`, it suffices to pass the `MnistModel`.\n",
        "Evaluation doesn't perform gradient descent, and there's no need to construct\n",
        "optimizers.\n",
        "\n",
        "For experimentation and research, when a centralized test dataset is available,\n",
        "[Federated Learning for Text Generation](federated_learning_for_text_generation.ipynb)\n",
        "demonstrates another evaluation option: taking the trained weights from\n",
        "federated learning, applying them to a standard Keras model, and then simply\n",
        "calling `tf.keras.models.Model.evaluate()` on a centralized dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "nRiXyqnXM2VO"
      },
      "outputs": [],
      "source": [
        "evaluation = tff.learning.build_federated_evaluation(MnistModel)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uwfINGoNQEuV"
      },
      "source": [
        "You can inspect the abstract type signature of the evaluation function as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 54
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 302,
          "status": "ok",
          "timestamp": 1584130596173,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "3q5ueoO0NDNb",
        "outputId": "9d73bb02-0aa0-4f29-e5c3-e3a67e0be766"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'(\u003c\u003ctrainable=\u003cfloat32[784,10],float32[10]\u003e,non_trainable=\u003c\u003e\u003e@SERVER,{\u003cx=float32[?,784],y=int32[?,1]\u003e*}@CLIENTS\u003e -\u003e \u003cnum_examples=float32@SERVER,loss=float32@SERVER,accuracy=float32@SERVER\u003e)'"
            ]
          },
          "execution_count": 35,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(evaluation.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "XA3v7f2SQs6q"
      },
      "source": [
        "No need to be concerned about the details at this point, just be aware that it\n",
        "takes the following general form, similar to `tff.templates.IterativeProcess.next`\n",
        "but with two important differences. First, we are not returning server state,\n",
        "since evaluation doesn't modify the model or any other aspect of state - you can\n",
        "think of it as stateless. Second, evaluation only needs the model, and doesn't\n",
        "require any other part of server state that might be associated with training,\n",
        "such as optimizer variables.\n",
        "\n",
        "```\n",
        "SERVER_MODEL, FEDERATED_DATA -\u003e TRAINING_METRICS\n",
        "```\n",
        "\n",
        "Let's invoke evaluation on the latest state we arrived at during training. In\n",
        "order to extract the latest trained model from the server state, you simply\n",
        "access the `.model` member, as follows."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OX4Sk_uyOaYa"
      },
      "outputs": [],
      "source": [
        "train_metrics = evaluation(state.model, federated_train_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UeEsdwJgRGMW"
      },
      "source": [
        "Here's what we get. Note the numbers look marginally better than what was\n",
        "reported by the last round of training above. By convention, the training\n",
        "metrics reported by the iterative training process generally reflect the\n",
        "performance of the model at the beginning of the training round, so the\n",
        "evaluation metrics will always be one step ahead."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 305,
          "status": "ok",
          "timestamp": 1584130598428,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "zwCy1IPxOfiT",
        "outputId": "b462d146-a050-46d1-ba7d-fa3ed1d62d2b"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'\u003cnum_examples=4860.0,loss=1.5626375675201416,accuracy=0.5133745074272156\u003e'"
            ]
          },
          "execution_count": 37,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(train_metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SpfgdNDoRjPy"
      },
      "source": [
        "Now, let's compile a test sample of federated data and rerun evaluation on the\n",
        "test data. The data will come from the same sample of real users, but from a\n",
        "distinct held-out data set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 51
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 475,
          "status": "ok",
          "timestamp": 1584130598910,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "in8vProVNc04",
        "outputId": "92a92fa4-61d3-40ea-c7cd-d1d7e7ba34f0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(10,\n",
              " \u003cDatasetV1Adapter shapes: OrderedDict([(x, (None, 784)), (y, (None, 1))]), types: OrderedDict([(x, tf.float32), (y, tf.int32)])\u003e)"
            ]
          },
          "execution_count": 38,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "federated_test_data = make_federated_data(emnist_test, sample_clients)\n",
        "\n",
        "len(federated_test_data), federated_test_data[0]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ty-ZwfE0NJfV"
      },
      "outputs": [],
      "source": [
        "test_metrics = evaluation(state.model, federated_test_data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "height": 34
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 307,
          "status": "ok",
          "timestamp": 1584130601231,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "e5fGtIJYNqYH",
        "outputId": "dc8f6f84-87bc-4911-a29a-39c86f19f204"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'\u003cnum_examples=580.0,loss=1.6954734325408936,accuracy=0.4568965435028076\u003e'"
            ]
          },
          "execution_count": 40,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(test_metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "67vYxrDWzRcj"
      },
      "source": [
        "This concludes the tutorial. We encourage you to play with the\n",
        "parameters (e.g., batch sizes, number of users, epochs, learning rates, etc.), to modify the code above to simulate training on random samples of users in\n",
        "each round, and to explore the other tutorials we've developed."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "Federated_Learning_for_Image_Classification.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}