docs/tutorials/custom_federated_algorithms_2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tqrD7Yzlmlsk"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2k8X1C1nmpKv"
      },
      "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": "32xflLc4NTx-"
      },
      "source": [
        "# Custom Federated Algorithms, Part 2: Implementing Federated Averaging"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jtATV6DlqPs0"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_2\"\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.82.0/docs/tutorials/custom_federated_algorithms_2.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.82.0/docs/tutorials/custom_federated_algorithms_2.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",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/custom_federated_algorithms_2.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_igJ2sfaNWS8"
      },
      "source": [
        "This tutorial is the second part of a two-part series that demonstrates how to\n",
        "implement custom types of federated algorithms in TFF using the\n",
        "[Federated Core (FC)](../federated_core.md), which serves as a foundation for\n",
        "the [Federated Learning (FL)](../federated_learning.md) layer (`tff.learning`).\n",
        "\n",
        "We encourage you to first read the\n",
        "[first part of this series](custom_federated_algorithms_1.ipynb), which\n",
        "introduce some of the key concepts and programming abstractions used here.\n",
        "\n",
        "This second part of the series uses the mechanisms introduced in the first part\n",
        "to implement a simple version of federated training and evaluation algorithms.\n",
        "\n",
        "We encourage you to review the\n",
        "[image classification](federated_learning_for_image_classification.ipynb) and\n",
        "[text generation](federated_learning_for_text_generation.ipynb) tutorials for a\n",
        "higher-level and more gentle introduction to TFF's Federated Learning APIs, as\n",
        "they will help you put the concepts we describe here in context."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cuJuLEh2TfZG"
      },
      "source": [
        "## Before we start\n",
        "\n",
        "Before we start, try to run the following \"Hello World\" example to make sure\n",
        "your environment is correctly setup. If it doesn't work, please refer to the\n",
        "[Installation](../install.md) guide for instructions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UqJnbs_1YKi4"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "-skNC6aovM46"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "zzXwGnZamIMM"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "b'Hello, World!'"
            ]
          },
          "execution_count": 2,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "@tff.federated_computation\n",
        "def hello_world():\n",
        "  return 'Hello, World!'\n",
        "\n",
        "hello_world()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iu5Gd8D6W33s"
      },
      "source": [
        "## Implementing Federated Averaging\n",
        "\n",
        "As in\n",
        "[Federated Learning for Image Classification](federated_learning_for_image_classification.ipynb),\n",
        "we are going to use the MNIST example, but since this is intended as a low-level\n",
        "tutorial, we are going to bypass the Keras API and `tff.simulation`, write raw\n",
        "model code, and construct a federated data set from scratch.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6qCjef350c_"
      },
      "source": [
        "### Preparing federated data sets\n",
        "\n",
        "For the sake of a demonstration, we're going to simulate a scenario in which we\n",
        "have data from 10 users, and each of the users contributes knowledge how to\n",
        "recognize a different digit. This is about as\n",
        "non-[i.i.d.](https://en.wikipedia.org/wiki/Independent_and_identically_distributed_random_variables)\n",
        "as it gets.\n",
        "\n",
        "First, let's load the standard MNIST data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "uThZM4Ds-KDQ"
      },
      "outputs": [],
      "source": [
        "mnist_train, mnist_test = tf.keras.datasets.mnist.load_data()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "PkJc5rHA2no_"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[(dtype('uint8'), (60000, 28, 28)), (dtype('uint8'), (60000,))]"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "[(x.dtype, x.shape) for x in mnist_train]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mFET4BKJFbkP"
      },
      "source": [
        "The data comes as Numpy arrays, one with images and another with digit labels, both\n",
        "with the first dimension going over the individual examples. Let's write a\n",
        "helper function that formats it in a way compatible with how we feed federated\n",
        "sequences into TFF computations, i.e., as a list of lists - the outer list\n",
        "ranging over the users (digits), the inner ones ranging over batches of data in\n",
        "each client's sequence. As is customary, we will structure each batch as a pair\n",
        "of tensors named `x` and `y`, each with the leading batch dimension. While at\n",
        "it, we'll also flatten each image into a 784-element vector and rescale the\n",
        "pixels in it into the `0..1` range, so that we don't have to clutter the model\n",
        "logic with data conversions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "XTaTLiq5GNqy"
      },
      "outputs": [],
      "source": [
        "NUM_EXAMPLES_PER_USER = 1000\n",
        "BATCH_SIZE = 100\n",
        "\n",
        "\n",
        "def get_data_for_digit(source, digit):\n",
        "  output_sequence = []\n",
        "  all_samples = [i for i, d in enumerate(source[1]) if d == digit]\n",
        "  for i in range(0, min(len(all_samples), NUM_EXAMPLES_PER_USER), BATCH_SIZE):\n",
        "    batch_samples = all_samples[i:i + BATCH_SIZE]\n",
        "    output_sequence.append({\n",
        "        'x':\n",
        "            np.array([source[0][i].flatten() / 255.0 for i in batch_samples],\n",
        "                     dtype=np.float32),\n",
        "        'y':\n",
        "            np.array([source[1][i] for i in batch_samples], dtype=np.int32)\n",
        "    })\n",
        "  return output_sequence\n",
        "\n",
        "\n",
        "federated_train_data = [get_data_for_digit(mnist_train, d) for d in range(10)]\n",
        "\n",
        "federated_test_data = [get_data_for_digit(mnist_test, d) for d in range(10)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xpNdBimWaMHD"
      },
      "source": [
        "As a quick sanity check, let's look at the `Y` tensor in the last batch of data\n",
        "contributed by the fifth client (the one corresponding to the digit `5`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "bTNuL1W4bcuc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "array([5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,\n",
              "       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,\n",
              "       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,\n",
              "       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5,\n",
              "       5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5], dtype=int32)"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "federated_train_data[5][-1]['y']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xgvcwv7Obhat"
      },
      "source": [
        "Just to be sure, let's also look at the image corresponding to the last element of that batch."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "cI4aat1za525"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90\nbGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAAsT\nAAALEwEAmpwYAAAN6klEQVR4nO3dfaxU9Z3H8c/Ha4sijQGMhlB20canjXGtEt2EZtHU1od/pBJI\nMTbqNqEJmlSzyS52/9Bk3WhcuutfPlAfYNdqUyNWggutAbN2MWm8GlaxbCurbotcQReM+BQVvvvH\nPWyueOc3l5kzcwa+71dyMzPne8853wz3wzkzvzPzc0QIwJHvqKYbANAfhB1IgrADSRB2IAnCDiRx\ndD93Zpu3/oEeiwiPt7yrI7vtS23/zvY228u62RaA3nKn4+y2hyT9XtK3JG2X9LykxRHx28I6HNmB\nHuvFkf18Sdsi4rWI+ETSzyRd0cX2APRQN2GfKemPYx5vr5Z9ju0ltodtD3exLwBd6uYNuvFOFb5w\nmh4RKyStkDiNB5rUzZF9u6RZYx5/VdKO7toB0CvdhP15SafaPtn2lyV9V9KaetoCULeOT+Mj4jPb\nN0j6paQhSQ9GxCu1dQagVh0PvXW0M16zAz3Xk4tqABw+CDuQBGEHkiDsQBKEHUiCsANJEHYgCcIO\nJEHYgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1Ioq9TNmPwXH/9\n9cX6Bx98UKyvXLmyxm4+b/bs2cX6UUeVj1WLFi1qWZs58wszlX3O0qVLi/WLL764WH/mmWeK9SZw\nZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnT27+/PnF+kUXXVSsT58+vVjfvHlzy9pVV11VXPfq\nq68u1oeGhor1brz//vvF+p49e3q2717pKuy235C0V9I+SZ9FxJw6mgJQvzqO7BdFxDs1bAdAD/Ga\nHUii27CHpF/ZfsH2kvF+wfYS28O2h7vcF4AudHsaPzcidtg+UdLTtv8rIp4d+wsRsULSCkmyHV3u\nD0CHujqyR8SO6naXpCcknV9HUwDq13HYbR9n+ysH7kv6tqQtdTUGoF7dnMafJOkJ2we280hErK+l\nKxw27rzzzmI9YjBfud10003F+rp164r1bdu21dlOX3Qc9oh4TdKf19gLgB5i6A1IgrADSRB2IAnC\nDiRB2IEk+IjrEaAa/hzX3Llzi+vOmzev7nYm7KOPPirW9+7dW6yvX18e6b3tttta1l5//fXiuoM6\nZNgNjuxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kIT7OZ7IN9X0xpQpU1rW3n333Z7u+5NPPinW16xZ\n07K2fPny4rrDw3yTWSciYtwLLziyA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EASfJ79CLBw4cLG9r10\n6dJifeXKlf1pBG1xZAeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJBhnPwwsWrSoWL/rrrt6tu+77767\nWGcc/fDR9shu+0Hbu2xvGbNsmu2nbb9a3U7tbZsAujWR0/iVki49aNkySRsi4lRJG6rHAAZY27BH\nxLOSdh+0+ApJq6r7qyTNr7ctAHXr9DX7SRExIkkRMWL7xFa/aHuJpCUd7gdATXr+Bl1ErJC0QuIL\nJ4EmdTr0ttP2DEmqbnfV1xKAXug07GskXVPdv0bSk/W0A6BX2n5vvO1HJV0o6QRJOyXdIukXkn4u\n6U8k/UHSwog4+E288bbFafw4Jk+eXKw/99xzxfpZZ53V8b43btxYrC9YsKBYbzeHOvqv1ffGt33N\nHhGLW5S+2VVHAPqKy2WBJAg7kARhB5Ig7EAShB1Igo+49sGkSZOK9fvuu69Y72ZorZ3bb7+9WGdo\n7cjBkR1IgrADSRB2IAnCDiRB2IEkCDuQBGEHkmCcvQ8uvPDCYn3x4lYfLOy9K6+8slg/++yzi/X3\n3nuvWH/ooYcOuSf0Bkd2IAnCDiRB2IEkCDuQBGEHkiDsQBKEHUii7VdJ17qzpF8l/dRTTxXrl156\n8LyZh4+jjiofL558svWUAu2elwceeKBY379/f7GeVauvkubIDiRB2IEkCDuQBGEHkiDsQBKEHUiC\nsANJMM7eB+eee26xfs899xTr5513Xsf73rp1a7E+MjJSrM+aNatYP+2004r1bv6+li1bVqwvX768\n420fyToeZ7f9oO1dtreMWXar7Tdtb65+Lq+zWQD1m8hp/EpJ413i9c8RcU7182/1tgWgbm3DHhHP\nStrdh14A9FA3b9DdYPul6jR/aqtfsr3E9rDt4S72BaBLnYb9Hklfk3SOpBFJP271ixGxIiLmRMSc\nDvcFoAYdhT0idkbEvojYL+knks6vty0Adeso7LZnjHn4HUlbWv0ugMHQdpzd9qOSLpR0gqSdkm6p\nHp8jKSS9IekHEVEesFXecfZ2Jk+eXKyfcsopHW/7zTffLNb37NlTrE+fPr1YP/3004v1m2++uWXt\nsssuK667b9++Yn3+/PnF+rp164r1I1Wrcfa2k0RExHgzGJS/VQDAwOFyWSAJwg4kQdiBJAg7kARh\nB5LgI641OPbYY4v1jz/+uFjv579Bvw0NDbWsbd68ubjumWeeWaxv2rSpWJ83b16xfqTiq6SB5Ag7\nkARhB5Ig7EAShB1IgrADSRB2IIm2n3rDqOOPP75l7ZFHHimuu3DhwmL9ww8/7Kinw8GUKVNa1o45\n5piutn300fz5HgqO7EAShB1IgrADSRB2IAnCDiRB2IEkCDuQBAOVEzRnTusJbS655JLiuu2mNW73\nue5BVhpHl6SHH364Ze3kk0+uux0UcGQHkiDsQBKEHUiCsANJEHYgCcIOJEHYgSQYZ++D9evXF+ul\naY0l6bHHHquznUNy7bXXFuu33HJLsT516tSO9/3pp58W6/fee2/H286o7ZHd9izbz9jeavsV2z+s\nlk+z/bTtV6vbzv9VAfTcRE7jP5P01xFxpqS/kHS97T+TtEzShog4VdKG6jGAAdU27BExEhEvVvf3\nStoqaaakKyStqn5tlaT5PeoRQA0O6TW77dmSvi7pN5JOiogRafQ/BNsntlhniaQlXfYJoEsTDrvt\nKZIel3RjRLxnjzt33BdExApJK6ptHLkzGAIDbkJDb7a/pNGg/zQiVleLd9qeUdVnSNrVmxYB1KHt\nlM0ePYSvkrQ7Im4cs/wfJf1vRNxhe5mkaRHxN222ddge2S+44IKWtY0bNxbXnTRpUt3tDIx2Z3il\nv689e/YU1203JHn//fcX61m1mrJ5IqfxcyV9T9LLtjdXy34k6Q5JP7f9fUl/kFT+cnQAjWob9oj4\nD0mt/vv+Zr3tAOgVLpcFkiDsQBKEHUiCsANJEHYgibbj7LXu7DAeZy+57rrrivV2H8UcGhqqs52+\najfO/vbbb7esLViwoLjupk2bOuopu1bj7BzZgSQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJxtn74Iwz\nzijWV69eXay3m/K5l9pNJ7127dpivXSNwVtvvdVJS2iDcXYgOcIOJEHYgSQIO5AEYQeSIOxAEoQd\nSIJxduAIwzg7kBxhB5Ig7EAShB1IgrADSRB2IAnCDiTRNuy2Z9l+xvZW26/Y/mG1/Fbbb9reXP1c\n3vt2AXSq7UU1tmdImhERL9r+iqQXJM2XtEjS+xGxfMI746IaoOdaXVQzkfnZRySNVPf32t4qaWa9\n7QHotUN6zW57tqSvS/pNtegG2y/ZftD21BbrLLE9bHu4u1YBdGPC18bbniLp3yX9Q0Sstn2SpHck\nhaS/1+ip/l+12Qan8UCPtTqNn1DYbX9J0lpJv4yIfxqnPlvS2og4q812CDvQYx1/EMaj03Q+IGnr\n2KBXb9wd8B1JW7ptEkDvTOTd+G9I+rWklyXtrxb/SNJiSedo9DT+DUk/qN7MK22LIzvQY12dxteF\nsAO9x+fZgeQIO5AEYQeSIOxAEoQdSIKwA0kQdiAJwg4kQdiBJAg7kARhB5Ig7EAShB1IgrADSbT9\nwsmavSPpf8Y8PqFaNogGtbdB7Uuit07V2duftir09fPsX9i5PRwRcxproGBQexvUviR661S/euM0\nHkiCsANJNB32FQ3vv2RQexvUviR661Rfemv0NTuA/mn6yA6gTwg7kEQjYbd9qe3f2d5me1kTPbRi\n+w3bL1fTUDc6P101h94u21vGLJtm+2nbr1a3486x11BvAzGNd2Ga8Uafu6anP+/7a3bbQ5J+L+lb\nkrZLel7S4oj4bV8bacH2G5LmRETjF2DY/ktJ70v6lwNTa9m+U9LuiLij+o9yakT87YD0dqsOcRrv\nHvXWaprxa9Xgc1fn9OedaOLIfr6kbRHxWkR8Iulnkq5ooI+BFxHPStp90OIrJK2q7q/S6B9L37Xo\nbSBExEhEvFjd3yvpwDTjjT53hb76oomwz5T0xzGPt2uw5nsPSb+y/YLtJU03M46TDkyzVd2e2HA/\nB2s7jXc/HTTN+MA8d51Mf96tJsI+3tQ0gzT+NzcizpV0maTrq9NVTMw9kr6m0TkARyT9uMlmqmnG\nH5d0Y0S812QvY43TV1+etybCvl3SrDGPvyppRwN9jCsidlS3uyQ9odGXHYNk54EZdKvbXQ338/8i\nYmdE7IuI/ZJ+ogafu2qa8ccl/TQiVleLG3/uxuurX89bE2F/XtKptk+2/WVJ35W0poE+vsD2cdUb\nJ7J9nKRva/Cmol4j6Zrq/jWSnmywl88ZlGm8W00zroafu8anP4+Ivv9Iulyj78j/t6S/a6KHFn2d\nIuk/q59Xmu5N0qMaPa37VKNnRN+XNF3SBkmvVrfTBqi3f9Xo1N4vaTRYMxrq7RsafWn4kqTN1c/l\nTT93hb768rxxuSyQBFfQAUkQdiAJwg4kQdiBJAg7kARhB5Ig7EAS/wfgQlrpjsiFUAAAAABJRU5E\nrkJggg==\n",
            "text/plain": [
              "\u003cFigure size 600x400 with 1 Axes\u003e"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "\n",
        "plt.imshow(federated_train_data[5][-1]['x'][-1].reshape(28, 28), cmap='gray')\n",
        "plt.grid(False)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J-ox58PA56f8"
      },
      "source": [
        "### On combining TensorFlow and TFF\n",
        "\n",
        "In this tutorial, for compactness we immediately decorate functions that\n",
        "introduce TensorFlow logic with `tff.tensorflow.computation`. However, for more complex\n",
        "logic, this is not the pattern we recommend. Debugging TensorFlow can already be\n",
        "a challenge, and debugging TensorFlow after it has been fully serialized and\n",
        "then re-imported necessarily loses some metadata and limits interactivity,\n",
        "making debugging even more of a challenge.\n",
        "\n",
        "Therefore, **we strongly recommend writing complex TF logic as stand-alone\n",
        "Python functions** (that is, without `tff.tensorflow.computation` decoration). This way\n",
        "the TensorFlow logic can be developed and tested using TF best practices and\n",
        "tools (like eager mode), before serializing the computation for TFF (e.g., by invoking `tff.tensorflow.computation` with a Python function as the argument)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RSd6UatXbzw-"
      },
      "source": [
        "### Defining a loss function\n",
        "\n",
        "Now that we have the data, let's define a loss function that we can use for\n",
        "training. First, let's define the type of input as a TFF named tuple. Since the\n",
        "size of data batches may vary, we set the batch dimension to `None` to indicate\n",
        "that the size of this dimension is unknown."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "653xv5NXd4fy"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'\u003cx=float32[?,784],y=int32[?]\u003e'"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "BATCH_SPEC = collections.OrderedDict(\n",
        "    x=tf.TensorSpec(shape=[None, 784], dtype=tf.float32),\n",
        "    y=tf.TensorSpec(shape=[None], dtype=tf.int32))\n",
        "BATCH_TYPE = tff.types.tensorflow_to_type(BATCH_SPEC)\n",
        "\n",
        "str(BATCH_TYPE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pb6qPUvyh5A1"
      },
      "source": [
        "You may be wondering why we can't just define an ordinary Python type. Recall\n",
        "the discussion in [part 1](custom_federated_algorithms_1.ipynb), where we\n",
        "explained that while we can express the logic of TFF computations using Python,\n",
        "under the hood TFF computations *are not* Python. The symbol `BATCH_TYPE`\n",
        "defined above represents an abstract TFF type specification. It is important to\n",
        "distinguish this *abstract* TFF type from concrete Python *representation*\n",
        "types, e.g., containers such as `dict` or `collections.namedtuple` that may be\n",
        "used to represent the TFF type in the body of a Python function. Unlike Python,\n",
        "TFF has a single abstract type constructor `tff.StructType` for tuple-like\n",
        "containers, with elements that can be individually named or left unnamed. This\n",
        "type is also used to model formal parameters of computations, as TFF\n",
        "computations can formally only declare one parameter and one result - you will\n",
        "see examples of this shortly.\n",
        "\n",
        "Let's now define the TFF type of model parameters, again as a TFF named tuple of\n",
        "*weights* and *bias*."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "Og7VViafh-30"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u003cweights=float32[784,10],bias=float32[10]\u003e\n"
          ]
        }
      ],
      "source": [
        "MODEL_SPEC = collections.OrderedDict(\n",
        "    weights=tf.TensorSpec(shape=[784, 10], dtype=tf.float32),\n",
        "    bias=tf.TensorSpec(shape=[10], dtype=tf.float32))\n",
        "MODEL_TYPE = tff.types.tensorflow_to_type(MODEL_SPEC)\n",
        "print(MODEL_TYPE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iHhdaWSpfQxo"
      },
      "source": [
        "With those definitions in place, now we can define the loss for a given model, over a single batch. Note the usage of `@tf.function` decorator inside the `@tff.tensorflow.computation` decorator. This allows us to write TF using Python like semantics even though were inside a `tf.Graph` context created by the `tff.tensorflow.computation` decorator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "4EObiz_Ke0uK"
      },
      "outputs": [],
      "source": [
        "# NOTE: `forward_pass` is defined separately from `batch_loss` so that it can \n",
        "# be later called from within another tf.function. Necessary because a\n",
        "# @tf.function  decorated method cannot invoke a @tff.tensorflow.computation.\n",
        "\n",
        "@tf.function\n",
        "def forward_pass(model, batch):\n",
        "  predicted_y = tf.nn.softmax(\n",
        "      tf.matmul(batch['x'], model['weights']) + model['bias'])\n",
        "  return -tf.reduce_mean(\n",
        "      tf.reduce_sum(\n",
        "          tf.one_hot(batch['y'], 10) * tf.math.log(predicted_y), axis=[1]))\n",
        "\n",
        "@tff.tensorflow.computation(MODEL_TYPE, BATCH_TYPE)\n",
        "def batch_loss(model, batch):\n",
        "  return forward_pass(model, batch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8K0UZHGnr8SB"
      },
      "source": [
        "As expected, computation `batch_loss` returns `float32` loss given the model and\n",
        "a single data batch. Note how the `MODEL_TYPE` and `BATCH_TYPE` have been lumped\n",
        "together into a 2-tuple of formal parameters; you can recognize the type of\n",
        "`batch_loss` as `(\u003cMODEL_TYPE,BATCH_TYPE\u003e -\u003e float32)`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "4WXEAY8Nr89V"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'(\u003cmodel=\u003cweights=float32[784,10],bias=float32[10]\u003e,batch=\u003cx=float32[?,784],y=int32[?]\u003e\u003e -\u003e float32)'"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(batch_loss.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pAnt_UcdnvGa"
      },
      "source": [
        "As a sanity check, let's construct an initial model filled with zeros and\n",
        "compute the loss over the batch of data we visualized above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "U8Ne8igan3os"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "2.3025851"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "initial_model = collections.OrderedDict(\n",
        "    weights=np.zeros([784, 10], dtype=np.float32),\n",
        "    bias=np.zeros([10], dtype=np.float32))\n",
        "\n",
        "sample_batch = federated_train_data[5][-1]\n",
        "\n",
        "batch_loss(initial_model, sample_batch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ckigEAyDAWFV"
      },
      "source": [
        "Note that we feed the TFF computation with the initial model defined as a\n",
        "`dict`, even though the body of the Python function that defines it consumes\n",
        "model parameters as `model['weight']` and `model['bias']`. The arguments of the call\n",
        "to `batch_loss` aren't simply passed to the body of that function.\n",
        "\n",
        "\n",
        "What happens when we invoke `batch_loss`?\n",
        "The Python body of `batch_loss` has already been traced and serialized  in the above cell where it was defined.  TFF acts as the caller to `batch_loss`\n",
        "at the computation definition time, and as the target of invocation at the time\n",
        "`batch_loss` is invoked. In both roles, TFF serves as the bridge between TFF's\n",
        "abstract type system and Python representation types. At the invocation time,\n",
        "TFF will accept most standard Python container types (`dict`, `list`, `tuple`,\n",
        "`collections.namedtuple`, etc.) as concrete representations of abstract TFF\n",
        "tuples. Also, although as noted above, TFF computations formally only accept a\n",
        "single parameter, you can use the familiar Python call syntax with positional\n",
        "and/or keyword arguments in case where the type of the parameter is a tuple - it\n",
        "works as expected."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eB510nILYbId"
      },
      "source": [
        "### Gradient descent on a single batch\n",
        "\n",
        "Now, let's define a computation that uses this loss function to perform a single\n",
        "step of gradient descent. Note how in defining this function, we use\n",
        "`batch_loss` as a subcomponent. You can invoke a computation constructed with\n",
        "`tff.tensorflow.computation` inside the body of another computation, though typically\n",
        "this is not necessary - as noted above, because serialization looses some\n",
        "debugging information, it is often preferable for more complex computations to\n",
        "write and test all the TensorFlow without the `tff.tensorflow.computation` decorator."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "O4uaVxw3AyYS"
      },
      "outputs": [],
      "source": [
        "@tff.tensorflow.computation(MODEL_TYPE, BATCH_TYPE, np.float32)\n",
        "def batch_train(initial_model, batch, learning_rate):\n",
        "  # Define a group of model variables and set them to `initial_model`. Must\n",
        "  # be defined outside the @tf.function.\n",
        "  model_vars = collections.OrderedDict([\n",
        "      (name, tf.Variable(name=name, initial_value=value))\n",
        "      for name, value in initial_model.items()\n",
        "  ])\n",
        "  optimizer = tf.keras.optimizers.SGD(learning_rate)\n",
        "\n",
        "  @tf.function\n",
        "  def _train_on_batch(model_vars, batch):\n",
        "    # Perform one step of gradient descent using loss from `batch_loss`.\n",
        "    with tf.GradientTape() as tape:\n",
        "      loss = forward_pass(model_vars, batch)\n",
        "    grads = tape.gradient(loss, model_vars)\n",
        "    optimizer.apply_gradients(\n",
        "        zip(tf.nest.flatten(grads), tf.nest.flatten(model_vars)))\n",
        "    return model_vars\n",
        "\n",
        "  return _train_on_batch(model_vars, batch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "id": "Y84gQsaohC38"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'(\u003cinitial_model=\u003cweights=float32[784,10],bias=float32[10]\u003e,batch=\u003cx=float32[?,784],y=int32[?]\u003e,learning_rate=float32\u003e -\u003e \u003cweights=float32[784,10],bias=float32[10]\u003e)'"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(batch_train.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ID8xg9FCUL2A"
      },
      "source": [
        "When you invoke a Python function decorated with `tff.tensorflow.computation` within the\n",
        "body of another such function, the logic of the inner TFF computation is\n",
        "embedded (essentially, inlined) in the logic of the outer one. As noted above,\n",
        "if you are writing both computations, it is likely preferable to make the inner\n",
        "function (`batch_loss` in this case) a regular Python or `tf.function` rather\n",
        "than a `tff.tensorflow.computation`. However, here we illustrate that calling one\n",
        "`tff.tensorflow.computation` inside another basically works as expected. This may be\n",
        "necessary if, for example, you do not have the Python code defining\n",
        "`batch_loss`, but only its serialized TFF representation.\n",
        "\n",
        "Now, let's apply this function a few times to the initial model to see whether\n",
        "the loss decreases."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "id": "8edcJTlXUULm"
      },
      "outputs": [],
      "source": [
        "model = initial_model\n",
        "losses = []\n",
        "for _ in range(5):\n",
        "  model = batch_train(model, sample_batch, 0.1)\n",
        "  losses.append(batch_loss(model, sample_batch))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "id": "3n1onojT1zHv"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[0.19690025, 0.13176318, 0.101132266, 0.08273812, 0.0703014]"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "losses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EQk4Ha8PU-3P"
      },
      "source": [
        "### Gradient descent on a sequence of local data\n",
        "\n",
        "Now, since `batch_train` appears to work, let's write a similar training\n",
        "function `local_train` that consumes the entire sequence of all batches from one\n",
        "user instead of just a single batch. The new computation will need to now\n",
        "consume `tff.SequenceType(BATCH_TYPE)` instead of `BATCH_TYPE`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "EfPD5a6QVNXM"
      },
      "outputs": [],
      "source": [
        "LOCAL_DATA_TYPE = tff.SequenceType(BATCH_TYPE)\n",
        "\n",
        "@tff.federated_computation(MODEL_TYPE, np.float32, LOCAL_DATA_TYPE)\n",
        "def local_train(initial_model, learning_rate, all_batches):\n",
        "  \n",
        "  # Reduction function to apply to each batch.\n",
        "  @tff.federated_computation((MODEL_TYPE, np.float32), BATCH_TYPE)\n",
        "  def batch_fn(model_with_lr, batch):\n",
        "    model, lr = model_with_lr\n",
        "    return batch_train(model, batch, lr), lr\n",
        "\n",
        "  trained_model, _ = tff.sequence_reduce(\n",
        "      all_batches, (initial_model, learning_rate), batch_fn\n",
        "  )\n",
        "  return trained_model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "sAhkS5yKUgjC"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'(\u003cinitial_model=\u003cweights=float32[784,10],bias=float32[10]\u003e,learning_rate=float32,all_batches=\u003cx=float32[?,784],y=int32[?]\u003e*\u003e -\u003e \u003cweights=float32[784,10],bias=float32[10]\u003e)'"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(local_train.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EYT-SiopYBtH"
      },
      "source": [
        "There are quite a few details buried in this short section of code, let's go\n",
        "over them one by one.\n",
        "\n",
        "First, while we could have implemented this logic entirely in TensorFlow,\n",
        "relying on `tf.data.Dataset.reduce` to process the sequence similarly to how\n",
        "we've done it earlier, we've opted this time to express the logic in the glue\n",
        "language, as a `tff.federated_computation`. We've used the federated operator\n",
        "`tff.sequence_reduce` to perform the reduction.\n",
        "\n",
        "The operator `tff.sequence_reduce` is used similarly to\n",
        "`tf.data.Dataset.reduce`. You can think of it as essentially the same as\n",
        "`tf.data.Dataset.reduce`, but for use inside federated computations, which as\n",
        "you may remember, cannot contain TensorFlow code. It is a template operator with\n",
        "a formal parameter 3-tuple that consists of a *sequence* of `T`-typed elements,\n",
        "the initial state of the reduction (we'll refer to it abstractly as *zero*) of\n",
        "some type `U`, and the *reduction operator* of type `(\u003cU,T\u003e -\u003e U)` that alters the\n",
        "state of the reduction by processing a single element. The result is the final\n",
        "state of the reduction, after processing all elements in a sequential order. In\n",
        "our example, the state of the reduction is the model trained on a prefix of the\n",
        "data, and the elements are data batches.\n",
        "\n",
        "Second, note that we have again used one computation (`batch_train`) as a\n",
        "component within another (`local_train`), but not directly. We can't use it as a\n",
        "reduction operator because it takes an additional parameter - the learning rate.\n",
        "To resolve this, we define an embedded federated computation `batch_fn` that\n",
        "binds to the `local_train`'s parameter `learning_rate` in its body. It is\n",
        "allowed for a child computation defined this way to capture a formal parameter\n",
        "of its parent as long as the child computation is not invoked outside the body\n",
        "of its parent. You can think of this pattern as an equivalent of\n",
        "`functools.partial` in Python.\n",
        "\n",
        "The practical implication of capturing `learning_rate` this way is, of course,\n",
        "that the same learning rate value is used across all batches.\n",
        "\n",
        "Now, let's try the newly defined local training function on the entire sequence\n",
        "of data from the same user who contributed the sample batch (digit `5`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "EnWFLoZGcSby"
      },
      "outputs": [],
      "source": [
        "locally_trained_model = local_train(initial_model, 0.1, federated_train_data[5])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0UXUqGk9zoF"
      },
      "source": [
        "Did it work? To answer this question, we need to implement evaluation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a8WDKu6WYy__"
      },
      "source": [
        "### Local evaluation\n",
        "\n",
        "Here's one way to implement local evaluation by adding up the losses across all data\n",
        "batches (we could have just as well computed the average; we'll leave it as an\n",
        "exercise for the reader)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "0RiODuc6z7Ln"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation(MODEL_TYPE, LOCAL_DATA_TYPE)\n",
        "def local_eval(model, all_batches):\n",
        "\n",
        "  @tff.tensorflow.computation((MODEL_TYPE, np.float32), BATCH_TYPE)\n",
        "  def accumulate_evaluation(model_and_accumulator, batch):\n",
        "    model, accumulator = model_and_accumulator\n",
        "    return model, accumulator + batch_loss(model, batch)\n",
        "\n",
        "  _, total_loss = tff.sequence_reduce(\n",
        "      all_batches, (model, 0.0), accumulate_evaluation\n",
        "  )\n",
        "  return total_loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "pH2XPEAKa4Dg"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'(\u003cmodel=\u003cweights=float32[784,10],bias=float32[10]\u003e,all_batches=\u003cx=float32[?,784],y=int32[?]\u003e*\u003e -\u003e float32)'"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "str(local_eval.type_signature)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efX81HuE-BcO"
      },
      "source": [
        "Again, there are a few new elements illustrated by this code, let's go over them\n",
        "one by one.\n",
        "\n",
        "First, we have used two new federated operators for processing sequences:\n",
        "`tff.sequence_map` that takes a *mapping function* `T-\u003eU` and a *sequence* of\n",
        "`T`, and emits a sequence of `U` obtained by applying the mapping function\n",
        "pointwise, and `tff.sequence_sum` that just adds all the elements. Here, we map\n",
        "each data batch to a loss value, and then add the resulting loss values to\n",
        "compute the total loss.\n",
        "\n",
        "Note that we could have again used `tff.sequence_reduce`, but this wouldn't be\n",
        "the best choice - the reduction process is, by definition, sequential, whereas\n",
        "the mapping and sum can be computed in parallel. When given a choice, it's best\n",
        "to stick with operators that don't constrain implementation choices, so that\n",
        "when our TFF computation is compiled in the future to be deployed to a specific\n",
        "environment, one can take full advantage of all potential opportunities for a\n",
        "faster, more scalable, more resource-efficient execution.\n",
        "\n",
        "Second, note that just as in `local_train`, the component function we need\n",
        "(`batch_loss`) takes more parameters than what the federated operator\n",
        "(`tff.sequence_map`) expects, so we again define a partial, this time inline by\n",
        "directly wrapping a `lambda` as a `tff.federated_computation`. Using wrappers\n",
        "inline with a function as an argument is the recommended way to use\n",
        "`tff.tensorflow.computation` to embed TensorFlow logic in TFF.\n",
        "\n",
        "Now, let's see whether our training worked."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "vPw6JSVf5q_x"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "initial_model loss = 23.025854\n",
            "locally_trained_model loss = 0.43484688\n"
          ]
        }
      ],
      "source": [
        "print('initial_model loss =', local_eval(initial_model,\n",
        "                                         federated_train_data[5]))\n",
        "print('locally_trained_model loss =',\n",
        "      local_eval(locally_trained_model, federated_train_data[5]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Tvu70cnBsUf"
      },
      "source": [
        "Indeed, the loss decreased. But what happens if we evaluated it on another\n",
        "user's data?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "gjF0NYAj5wls"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "initial_model loss = 23.025854\n",
            "locally_trained_model loss = 74.50075\n"
          ]
        }
      ],
      "source": [
        "print('initial_model loss =', local_eval(initial_model,\n",
        "                                         federated_train_data[0]))\n",
        "print('locally_trained_model loss =',\n",
        "      local_eval(locally_trained_model, federated_train_data[0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7WPumnRTBzUs"
      },
      "source": [
        "As expected, things got worse. The model was trained to recognize `5`, and has\n",
        "never seen a `0`. This brings the question - how did the local training impact\n",
        "the quality of the model from the global perspective?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QJnL2mQRZKTO"
      },
      "source": [
        "### Federated evaluation\n",
        "\n",
        "This is the point in our journey where we finally circle back to federated types\n",
        "and federated computations - the topic that we started with. Here's a pair of\n",
        "TFF types definitions for the model that originates at the server, and the data\n",
        "that remains on the clients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "id": "LjGGhpoEBh_6"
      },
      "outputs": [],
      "source": [
        "SERVER_MODEL_TYPE = tff.FederatedType(MODEL_TYPE, tff.SERVER)\n",
        "CLIENT_DATA_TYPE = tff.FederatedType(LOCAL_DATA_TYPE, tff.CLIENTS)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4gTXV2-jZtE3"
      },
      "source": [
        "With all the definitions introduced so far, expressing federated evaluation in\n",
        "TFF is a one-liner - we distribute the model to clients, let each client invoke\n",
        "local evaluation on its local portion of data, and then average out the loss.\n",
        "Here's one way to write this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "id": "2zChEPzEBx4T"
      },
      "outputs": [],
      "source": [
        "@tff.federated_computation(SERVER_MODEL_TYPE, CLIENT_DATA_TYPE)\n",
        "def federated_eval(model, data):\n",
        "  return tff.federated_mean(\n",
        "      tff.federated_map(local_eval, [tff.federated_broadcast(model),  data]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IWcNONNWaE0N"
      },
      "source": [
        "We've already seen examples of `tff.federated_mean` and `tff.federated_map`\n",
        "in simpler scenarios, and at the intuitive level, they work as expected, but\n",
        "there's more in this section of code than meets the eye, so let's go over it\n",
        "carefully.\n",
        "\n",
        "First, let's break down the *let each client invoke local evaluation on its\n",
        "local portion of data* part. As you may recall from the preceding sections,\n",
        "`local_eval` has a type signature of the form `(\u003cMODEL_TYPE, LOCAL_DATA_TYPE\u003e -\u003e\n",
        "float32)`.\n",
        "\n",
        "The federated operator `tff.federated_map` is a template that accepts as a\n",
        "parameter a 2-tuple that consists of the *mapping function* of some type `T-\u003eU`\n",
        "and a federated value of type `{T}@CLIENTS` (i.e., with member constituents of\n",
        "the same type as the parameter of the mapping function), and returns a result of\n",
        "type `{U}@CLIENTS`.\n",
        "\n",
        "Since we're feeding `local_eval` as a mapping function to apply on a per-client\n",
        "basis, the second argument should be of a federated type `{\u003cMODEL_TYPE,\n",
        "LOCAL_DATA_TYPE\u003e}@CLIENTS`, i.e., in the nomenclature of the preceding sections,\n",
        "it should be a federated tuple. Each client should hold a full set of arguments\n",
        "for `local_eval` as a member consituent. Instead, we're feeding it a 2-element\n",
        "Python `list`. What's happening here?\n",
        "\n",
        "Indeed, this is an example of an *implicit type cast* in TFF, similar to\n",
        "implicit type casts you may have encountered elsewhere, e.g., when you feed an\n",
        "`int` to a function that accepts a `float`. Implicit casting is used scarcily at\n",
        "this point, but we plan to make it more pervasive in TFF as a way to minimize\n",
        "boilerplate.\n",
        "\n",
        "The implicit cast that's applied in this case is the equivalence between\n",
        "federated tuples of the form `{\u003cX,Y\u003e}@Z`, and tuples of federated values\n",
        "`\u003c{X}@Z,{Y}@Z\u003e`. While formally, these two are different type signatures,\n",
        "looking at it from the programmers's perspective, each device in `Z` holds two\n",
        "units of data `X` and `Y`. What happens here is not unlike `zip` in Python, and\n",
        "indeed, we offer an operator `tff.federated_zip` that allows you to perform such\n",
        "conversions explicity. When the `tff.federated_map` encounters a tuple as a\n",
        "second argument, it simply invokes `tff.federated_zip` for you.\n",
        "\n",
        "Given the above, you should now be able to recognize the expression\n",
        "`tff.federated_broadcast(model)` as representing a value of TFF type\n",
        "`{MODEL_TYPE}@CLIENTS`, and `data` as a value of TFF type\n",
        "`{LOCAL_DATA_TYPE}@CLIENTS` (or simply `CLIENT_DATA_TYPE`), the two getting\n",
        "filtered together through an implicit `tff.federated_zip` to form the second\n",
        "argument to `tff.federated_map`.\n",
        "\n",
        "The operator `tff.federated_broadcast`, as you'd expect, simply transfers data\n",
        "from the server to the clients.\n",
        "\n",
        "Now, let's see how our local training affected the average loss in the system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "id": "tbmtJItcn94j"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "initial_model loss = 23.025852\n",
            "locally_trained_model loss = 54.43263\n"
          ]
        }
      ],
      "source": [
        "print('initial_model loss =', federated_eval(initial_model,\n",
        "                                             federated_train_data))\n",
        "print('locally_trained_model loss =',\n",
        "      federated_eval(locally_trained_model, federated_train_data))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LQi2rGX_fK7i"
      },
      "source": [
        "Indeed, as expected, the loss has increased. In order to improve the model for\n",
        "all users, we'll need to train in on everyone's data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkw9f59qfS7o"
      },
      "source": [
        "### Federated training\n",
        "\n",
        "The simplest way to implement federated training is to locally train, and then\n",
        "average the models. This uses the same building blocks and patters we've already\n",
        "discussed, as you can see below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "mBOC4uoG6dd-"
      },
      "outputs": [],
      "source": [
        "SERVER_FLOAT_TYPE = tff.FederatedType(np.float32, tff.SERVER)\n",
        "\n",
        "\n",
        "@tff.federated_computation(SERVER_MODEL_TYPE, SERVER_FLOAT_TYPE,\n",
        "                           CLIENT_DATA_TYPE)\n",
        "def federated_train(model, learning_rate, data):\n",
        "  return tff.federated_mean(\n",
        "      tff.federated_map(local_train, [\n",
        "          tff.federated_broadcast(model),\n",
        "          tff.federated_broadcast(learning_rate), data\n",
        "      ]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z2vACMsQjzO1"
      },
      "source": [
        "Note that in the full-featured implementation of Federated Averaging provided by\n",
        "`tff.learning`, rather than averaging the models, we prefer to average model\n",
        "deltas, for a number of reasons, e.g., the ability to clip the update norms,\n",
        "for compression, etc.\n",
        "\n",
        "Let's see whether the training works by running a few rounds of training and\n",
        "comparing the average loss before and after."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "NLx-3rLs9jGY"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round 0, loss=21.60552406311035\n",
            "round 1, loss=20.365678787231445\n",
            "round 2, loss=19.27480125427246\n",
            "round 3, loss=18.31110954284668\n",
            "round 4, loss=17.457256317138672\n"
          ]
        }
      ],
      "source": [
        "model = initial_model\n",
        "learning_rate = 0.1\n",
        "for round_num in range(5):\n",
        "  model = federated_train(model, learning_rate, federated_train_data)\n",
        "  learning_rate = learning_rate * 0.9\n",
        "  loss = federated_eval(model, federated_train_data)\n",
        "  print('round {}, loss={}'.format(round_num, loss))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0VjSLQzlUIp"
      },
      "source": [
        "For completeness, let's now also run on the test data to confirm that our model\n",
        "generalizes well."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "id": "ZaZT45yFMOaM"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "initial_model test loss = 22.795593\n",
            "trained_model test loss = 17.278767\n"
          ]
        }
      ],
      "source": [
        "print('initial_model test loss =',\n",
        "      federated_eval(initial_model, federated_test_data))\n",
        "print('trained_model test loss =', federated_eval(model, federated_test_data))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxlHHwLGlgFB"
      },
      "source": [
        "This concludes our tutorial.\n",
        "\n",
        "Of course, our simplified example doesn't reflect a number of things you'd need\n",
        "to do in a more realistic scenario - for example, we haven't computed metrics\n",
        "other than loss. We encourage you to study\n",
        "[the implementation](https://github.com/tensorflow/federated/blob/main/tensorflow_federated/python/learning/federated_averaging.py)\n",
        "of federated averaging in `tff.learning` as a more complete example, and as a\n",
        "way to demonstrate some of the coding practices we'd like to encourage."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "custom_federated_algorithms_2.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}