docs/tutorials/simulations_with_accelerators.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1YRGLD1pOOrJ"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Federated Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "koW3R4ntOgLS"
      },
      "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": "grBmytrShbUE"
      },
      "source": [
        "# TFF simulations with accelerators\n",
        "\n",
        "This tutorial will describe how to setup TFF simulations with accelerators. We focus on single-machine (multi-)GPU for now and will update this tutorial with multi-machine and TPU settings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "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/simulations_with_accelerators\"\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.77.0/docs/tutorials/simulations_with_accelerators.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.77.0/docs/tutorials/simulations_with_accelerators.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/simulations_with_accelerators.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": "yiq_MY4LopET"
      },
      "source": [
        "## Before we begin\n",
        "\n",
        "First, let us make sure the notebook is connected to a backend that has the relevant components compiled. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ke7EyuvG0Zyn"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated\n",
        "!pip install -U tensorboard_plugin_profile"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CLMBbjjGdG3W"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rtgStTrNIId-"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0FZdImo7QfkT"
      },
      "source": [
        "Check if TF can detect physical GPUs and create a virtual multi-GPU environment for TFF GPU simulations. The two virtual GPUs will have limited memory to demonstrate how to configure TFF runtime. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UpuzBJmQIw-J",
        "outputId": "4efba639-de66-4d27-a7fa-a4e1249f9ba6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[LogicalDevice(name='/device:CPU:0', device_type='CPU'),\n",
              " LogicalDevice(name='/device:GPU:0', device_type='GPU'),\n",
              " LogicalDevice(name='/device:GPU:1', device_type='GPU')]"
            ]
          },
          "execution_count": 2,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "gpu_devices = tf.config.list_physical_devices('GPU')\n",
        "if not gpu_devices:\n",
        "  raise ValueError('Cannot detect physical GPU device in TF')\n",
        "# TODO: b/277213652 - Remove this call, as it doesn't work with C++ executor\n",
        "tf.config.set_logical_device_configuration(\n",
        "    gpu_devices[0], \n",
        "    [tf.config.LogicalDeviceConfiguration(memory_limit=1024),\n",
        "     tf.config.LogicalDeviceConfiguration(memory_limit=1024)])\n",
        "tf.config.list_logical_devices()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X6eWsahmQpmi"
      },
      "source": [
        "Run the following \"Hello World\"\n",
        "example to make sure the TFF environment is correctly setup. If it doesn't work,\n",
        "please refer to the [Installation](../install.md) guide for instructions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wjX3wmC-P1aE",
        "outputId": "54bb0a3a-8104-4248-edc6-d7edcfef2771"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "b'Hello, World!'"
            ]
          },
          "execution_count": 3,
          "metadata": {
            "tags": []
          },
          "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": "C2-BdlyAId1_"
      },
      "source": [
        "## EMNIST experimental setup\n",
        "\n",
        "In this tutorial, we train an EMNIST image classifier with Federated Averaging algorithm. Let us start by loading  the MNIST example from the TFF website."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2dVPgxN0MdG2"
      },
      "outputs": [],
      "source": [
        "emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data(only_digits=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QJZh5NoiG-hR"
      },
      "source": [
        "We define a function preprocessing the EMNIST example following the [simple_fedavg](https://github.com/tensorflow/federated/tree/main/examples/simple_fedavg) example. Note that the argument `client_epochs_per_round` controls the number of local epochs on clients in federated learning. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0S7t0-3hHCWc"
      },
      "outputs": [],
      "source": [
        "def preprocess_emnist_dataset(client_epochs_per_round, batch_size, test_batch_size):\n",
        "\n",
        "  def element_fn(element):\n",
        "    return collections.OrderedDict(\n",
        "        x=tf.expand_dims(element['pixels'], -1), y=element['label'])\n",
        "\n",
        "  def preprocess_train_dataset(dataset):\n",
        "    # Use buffer_size same as the maximum client dataset size,\n",
        "    # 418 for Federated EMNIST\n",
        "    return dataset.map(element_fn).shuffle(buffer_size=418).repeat(\n",
        "        count=client_epochs_per_round).batch(batch_size, drop_remainder=False)\n",
        "\n",
        "  def preprocess_test_dataset(dataset):\n",
        "    return dataset.map(element_fn).batch(test_batch_size, drop_remainder=False)\n",
        "\n",
        "  train_set = emnist_train.preprocess(preprocess_train_dataset)\n",
        "  test_set = preprocess_test_dataset(\n",
        "      emnist_test.create_tf_dataset_from_all_clients())\n",
        "  return train_set, test_set"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6vmWv0ALKvqh"
      },
      "source": [
        "We use a VGG-like model, i.e., each block has two 3x3 convolutions and number of filters are doubled when the feature maps are subsampled. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D_5T0UzHKtde"
      },
      "outputs": [],
      "source": [
        "def _conv_3x3(input_tensor, filters, strides):\n",
        "  \"\"\"2D Convolutional layer with kernel size 3x3.\"\"\"\n",
        "\n",
        "  x = tf.keras.layers.Conv2D(\n",
        "      filters=filters,\n",
        "      strides=strides,\n",
        "      kernel_size=3,\n",
        "      padding='same',\n",
        "      kernel_initializer='he_normal',\n",
        "      use_bias=False,\n",
        "  )(input_tensor)\n",
        "  return x\n",
        "\n",
        "\n",
        "def _basic_block(input_tensor, filters, strides):\n",
        "  \"\"\"A block of two 3x3 conv layers.\"\"\"\n",
        "\n",
        "  x = input_tensor\n",
        "  x = _conv_3x3(x, filters, strides)\n",
        "  x = tf.keras.layers.Activation('relu')(x)\n",
        "\n",
        "  x = _conv_3x3(x, filters, 1)\n",
        "  x = tf.keras.layers.Activation('relu')(x)\n",
        "  return x\n",
        "\n",
        "\n",
        "def _vgg_block(input_tensor, size, filters, strides):\n",
        "  \"\"\"A stack of basic blocks.\"\"\"\n",
        "  x = _basic_block(input_tensor, filters, strides=strides)\n",
        "  for _ in range(size - 1):\n",
        "      x = _basic_block(x, filters, strides=1)\n",
        "  return x\n",
        "\n",
        "\n",
        "def create_cnn(num_blocks, conv_width_multiplier=1, num_classes=10):\n",
        "  \"\"\"Create a VGG-like CNN model. \n",
        "  \n",
        "  The CNN has (6*num_blocks + 2) layers.\n",
        "  \"\"\"\n",
        "  input_shape = (28, 28, 1)  # channels_last\n",
        "  img_input = tf.keras.layers.Input(shape=input_shape)\n",
        "  x = img_input\n",
        "  x = tf.image.per_image_standardization(x)\n",
        "\n",
        "  x = _conv_3x3(x, 16 * conv_width_multiplier, 1)\n",
        "  x = _vgg_block(x, size=num_blocks, filters=16 * conv_width_multiplier, strides=1)\n",
        "  x = _vgg_block(x, size=num_blocks, filters=32 * conv_width_multiplier, strides=2)\n",
        "  x = _vgg_block(x, size=num_blocks, filters=64 * conv_width_multiplier, strides=2)\n",
        "\n",
        "  x = tf.keras.layers.GlobalAveragePooling2D()(x)\n",
        "  x = tf.keras.layers.Dense(num_classes)(x)\n",
        "\n",
        "  model = tf.keras.models.Model(\n",
        "      img_input,\n",
        "      x,\n",
        "      name='cnn-{}-{}'.format(6 * num_blocks + 2, conv_width_multiplier))\n",
        "  return model\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8OGBU1SLpyiR"
      },
      "source": [
        "Now let us define the training loop for EMNIST. Note that `use_experimental_simulation_loop=True` in `tff.learning.algorithms.build_weighted_fed_avg` is suggested for performant TFF simulation, and required to take advantage of multi-GPUs on a single machine. See [simple_fedavg](https://github.com/tensorflow/federated/tree/main/examples/simple_fedavg) example for how to define customized federated learning algorithm that has high performance on GPUs, one of the key features is to explicitly use `for ... iter(dataset)` for training loops. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dh_rQ6qQEXNn"
      },
      "outputs": [],
      "source": [
        "def keras_evaluate(model, test_data, metric):\n",
        "  metric.reset_states()\n",
        "  for batch in test_data:\n",
        "    preds = model(batch['x'], training=False)\n",
        "    metric.update_state(y_true=batch['y'], y_pred=preds)\n",
        "  return metric.result()\n",
        "\n",
        "\n",
        "def run_federated_training(client_epochs_per_round, \n",
        "                           train_batch_size, \n",
        "                           test_batch_size, \n",
        "                           cnn_num_blocks, \n",
        "                           conv_width_multiplier,\n",
        "                           server_learning_rate, \n",
        "                           client_learning_rate, \n",
        "                           total_rounds, \n",
        "                           clients_per_round, \n",
        "                           rounds_per_eval,\n",
        "                           logdir='logdir'):\n",
        "\n",
        "  train_data, test_data = preprocess_emnist_dataset(\n",
        "      client_epochs_per_round, train_batch_size, test_batch_size)\n",
        "  data_spec = test_data.element_spec\n",
        "\n",
        "  def _model_fn():\n",
        "    keras_model = create_cnn(cnn_num_blocks, conv_width_multiplier)\n",
        "    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "    return tff.learning.models.from_keras_model(\n",
        "        keras_model, input_spec=data_spec, loss=loss)\n",
        "  \n",
        "  def _server_optimizer_fn():\n",
        "    return tf.keras.optimizers.SGD(learning_rate=server_learning_rate)\n",
        "\n",
        "  def _client_optimizer_fn():\n",
        "    return tf.keras.optimizers.SGD(learning_rate=client_learning_rate)\n",
        "    \n",
        "  learning_process = tff.learning.algorithms.build_weighted_fed_avg(\n",
        "      model_fn=_model_fn, \n",
        "      server_optimizer_fn=_server_optimizer_fn, \n",
        "      client_optimizer_fn=_client_optimizer_fn, \n",
        "      use_experimental_simulation_loop=True)\n",
        "  \n",
        "  metric = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')\n",
        "  eval_model = create_cnn(cnn_num_blocks, conv_width_multiplier)\n",
        "  logging.info(eval_model.summary())\n",
        "\n",
        "  server_state = learning_process.initialize()\n",
        "  start_time = time.time()\n",
        "  for round_num in range(total_rounds):\n",
        "    sampled_clients = np.random.choice(\n",
        "        train_data.client_ids,\n",
        "        size=clients_per_round,\n",
        "        replace=False)\n",
        "    sampled_train_data = [\n",
        "        train_data.create_tf_dataset_for_client(client)\n",
        "        for client in sampled_clients\n",
        "    ]\n",
        "    if round_num == total_rounds-1:\n",
        "      with tf.profiler.experimental.Profile(logdir):\n",
        "        result = learning_process.next(\n",
        "            server_state, sampled_train_data)\n",
        "    else:\n",
        "      result = learning_process.next(\n",
        "            server_state, sampled_train_data)\n",
        "    server_state = result.state\n",
        "    train_metrics = result.metrics['client_work']['train']\n",
        "    print(f'Round {round_num} training loss: {train_metrics[\"loss\"]}, '\n",
        "     f'time: {(time.time()-start_time)/(round_num+1.)} secs')\n",
        "    if round_num % rounds_per_eval == 0 or round_num == total_rounds-1:\n",
        "      model_weights = learning_process.get_model_weights(server_state)\n",
        "      model_weights.assign_weights_to(eval_model)\n",
        "      accuracy = keras_evaluate(eval_model, test_data, metric)\n",
        "      print(f'Round {round_num} validation accuracy: {accuracy * 100.0}')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r_IqGFKbro1v"
      },
      "source": [
        "## Single GPU execution\n",
        "The default runtime of TFF is the same as TF: when GPUs are provided, the first GPU will be chosen for execution. We run the previously defined federated training for several rounds with a relatively small model. The last round of execution is profiled with `tf.profiler` and visualized by `tensorboard`. The profiling verified the first GPU is used.  \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dyDMZpaUInOH",
        "outputId": "605777b3-34f5-4aca-9d9d-54b818302bf4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"cnn-14-4\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "input_1 (InputLayer)         [(None, 28, 28, 1)]       0         \n",
            "_________________________________________________________________\n",
            "tf.image.per_image_standardi (None, 28, 28, 1)         0         \n",
            "_________________________________________________________________\n",
            "conv2d (Conv2D)              (None, 28, 28, 64)        576       \n",
            "_________________________________________________________________\n",
            "conv2d_1 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation (Activation)      (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_2 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_1 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_3 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_2 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_4 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_3 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_5 (Conv2D)            (None, 14, 14, 128)       73728     \n",
            "_________________________________________________________________\n",
            "activation_4 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_6 (Conv2D)            (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_5 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_7 (Conv2D)            (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_6 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_8 (Conv2D)            (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_7 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_9 (Conv2D)            (None, 7, 7, 256)         294912    \n",
            "_________________________________________________________________\n",
            "activation_8 (Activation)    (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_10 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_9 (Activation)    (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_11 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_10 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_12 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_11 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "global_average_pooling2d (Gl (None, 256)               0         \n",
            "_________________________________________________________________\n",
            "dense (Dense)                (None, 10)                2570      \n",
            "=================================================================\n",
            "Total params: 2,731,082\n",
            "Trainable params: 2,731,082\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "Round 0 training loss: 2.4688243865966797, time: 13.382015466690063 secs\n",
            "Round 0 validation accuracy: 15.240497589111328\n",
            "Round 1 training loss: 2.3217368125915527, time: 9.311999917030334 secs\n",
            "Round 2 training loss: 2.3100595474243164, time: 6.972411632537842 secs\n",
            "Round 2 validation accuracy: 11.226489067077637\n",
            "Round 3 training loss: 2.303222417831421, time: 6.467299699783325 secs\n",
            "Round 4 training loss: 2.2976326942443848, time: 5.526083135604859 secs\n",
            "Round 4 validation accuracy: 11.224040031433105\n",
            "Round 5 training loss: 2.2919719219207764, time: 5.468692660331726 secs\n",
            "Round 6 training loss: 2.2911534309387207, time: 4.935825347900391 secs\n",
            "Round 6 validation accuracy: 11.833855628967285\n",
            "Round 7 training loss: 2.2871201038360596, time: 4.918408691883087 secs\n",
            "Round 8 training loss: 2.2818832397460938, time: 4.602836343977186 secs\n",
            "Round 8 validation accuracy: 11.385677337646484\n",
            "Round 9 training loss: 2.2790346145629883, time: 4.99558527469635 secs\n",
            "Round 9 validation accuracy: 11.226489067077637\n"
          ]
        }
      ],
      "source": [
        "run_federated_training(\n",
        "    client_epochs_per_round=1, \n",
        "    train_batch_size=16, \n",
        "    test_batch_size=128, \n",
        "    cnn_num_blocks=2, \n",
        "    conv_width_multiplier=4,\n",
        "    server_learning_rate=1.0, \n",
        "    client_learning_rate=0.01,\n",
        "    total_rounds=10,\n",
        "    clients_per_round=16, \n",
        "    rounds_per_eval=2,\n",
        "    )\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oQby1JNRfnby"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "%tensorboard --logdir=logdir --port=0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mZsODPDX1lEv"
      },
      "source": [
        "## Larger model and OOM\n",
        "Let us run a larger model on CPU with less federated rounds. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "muYz4b0-1usE",
        "outputId": "81c247d2-baba-42fe-aab6-e4e5bfaee09e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"cnn-26-4\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "input_4 (InputLayer)         [(None, 28, 28, 1)]       0         \n",
            "_________________________________________________________________\n",
            "tf.image.per_image_standardi (None, 28, 28, 1)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_39 (Conv2D)           (None, 28, 28, 64)        576       \n",
            "_________________________________________________________________\n",
            "conv2d_40 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_36 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_41 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_37 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_42 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_38 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_43 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_39 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_44 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_40 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_45 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_41 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_46 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_42 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_47 (Conv2D)           (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_43 (Activation)   (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_48 (Conv2D)           (None, 14, 14, 128)       73728     \n",
            "_________________________________________________________________\n",
            "activation_44 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_49 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_45 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_50 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_46 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_51 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_47 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_52 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_48 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_53 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_49 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_54 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_50 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_55 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_51 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_56 (Conv2D)           (None, 7, 7, 256)         294912    \n",
            "_________________________________________________________________\n",
            "activation_52 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_57 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_53 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_58 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_54 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_59 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_55 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_60 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_56 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_61 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_57 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_62 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_58 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_63 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_59 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "global_average_pooling2d_3 ( (None, 256)               0         \n",
            "_________________________________________________________________\n",
            "dense_3 (Dense)              (None, 10)                2570      \n",
            "=================================================================\n",
            "Total params: 5,827,658\n",
            "Trainable params: 5,827,658\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "Round 0 training loss: 2.437223434448242, time: 24.121686458587646 secs\n",
            "Round 0 validation accuracy: 9.024785041809082\n",
            "Round 1 training loss: 2.3081459999084473, time: 19.48685622215271 secs\n",
            "Round 2 training loss: 2.305305242538452, time: 15.73950457572937 secs\n",
            "Round 2 validation accuracy: 9.791339874267578\n",
            "Round 3 training loss: 2.303149700164795, time: 15.194068729877472 secs\n",
            "Round 4 training loss: 2.3026506900787354, time: 14.036769819259643 secs\n",
            "Round 4 validation accuracy: 12.193867683410645\n"
          ]
        }
      ],
      "source": [
        "run_federated_training(\n",
        "    client_epochs_per_round=1, \n",
        "    train_batch_size=16, \n",
        "    test_batch_size=128, \n",
        "    cnn_num_blocks=4, \n",
        "    conv_width_multiplier=4,\n",
        "    server_learning_rate=1.0, \n",
        "    client_learning_rate=0.01,\n",
        "    total_rounds=5,\n",
        "    clients_per_round=16, \n",
        "    rounds_per_eval=2,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FVPGSnSaBt8V"
      },
      "source": [
        "This model might hit an out of memory issue on a single GPU. The migration from large scale CPU experiments to GPU simulation can be constrained by memory usage as GPUs often have limited memeories. There are several parameters can be tuned in TFF runtime to mitigate OOM issue\n",
        "\n",
        "\n",
        "*   Adjust `max_concurrent_computation_calls` in `tff.backends.native.set_sync_local_cpp_execution_context` to control the concurrency of client traininng. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QhdtvKO8Dy5j",
        "outputId": "b95e05dc-81d4-4887-f697-3e5099ef1e90"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"cnn-26-4\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "input_1 (InputLayer)         [(None, 28, 28, 1)]       0         \n",
            "_________________________________________________________________\n",
            "tf.image.per_image_standardi (None, 28, 28, 1)         0         \n",
            "_________________________________________________________________\n",
            "conv2d (Conv2D)              (None, 28, 28, 64)        576       \n",
            "_________________________________________________________________\n",
            "conv2d_1 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation (Activation)      (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_2 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_1 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_3 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_2 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_4 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_3 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_5 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_4 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_6 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_5 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_7 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_6 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_8 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_7 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_9 (Conv2D)            (None, 14, 14, 128)       73728     \n",
            "_________________________________________________________________\n",
            "activation_8 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_10 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_9 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_11 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_10 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_12 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_11 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_13 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_12 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_14 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_13 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_15 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_14 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_16 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_15 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_17 (Conv2D)           (None, 7, 7, 256)         294912    \n",
            "_________________________________________________________________\n",
            "activation_16 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_18 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_17 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_19 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_18 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_20 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_19 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_21 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_20 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_22 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_21 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_23 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_22 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_24 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_23 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "global_average_pooling2d (Gl (None, 256)               0         \n",
            "_________________________________________________________________\n",
            "dense (Dense)                (None, 10)                2570      \n",
            "=================================================================\n",
            "Total params: 5,827,658\n",
            "Trainable params: 5,827,658\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "Round 0 training loss: 2.4990053176879883, time: 11.922378778457642 secs\n",
            "Round 0 validation accuracy: 11.224040031433105\n",
            "Round 1 training loss: 2.307560920715332, time: 9.916815996170044 secs\n",
            "Round 2 training loss: 2.3032877445220947, time: 7.68927804629008 secs\n",
            "Round 2 validation accuracy: 11.224040031433105\n",
            "Round 3 training loss: 2.302366256713867, time: 7.681552231311798 secs\n",
            "Round 4 training loss: 2.301671028137207, time: 7.613566827774048 secs\n",
            "Round 4 validation accuracy: 11.224040031433105\n"
          ]
        }
      ],
      "source": [
        "# Control concurrency by `max_concurrent_computation_calls`.\n",
        "tff.backends.native.set_sync_local_cpp_execution_context(\n",
        "    max_concurrent_computation_calls=16/2)\n",
        "\n",
        "run_federated_training(\n",
        "    client_epochs_per_round=1, \n",
        "    train_batch_size=16, \n",
        "    test_batch_size=128, \n",
        "    cnn_num_blocks=4, \n",
        "    conv_width_multiplier=4,\n",
        "    server_learning_rate=1.0, \n",
        "    client_learning_rate=0.01,\n",
        "    total_rounds=5,\n",
        "    clients_per_round=16, \n",
        "    rounds_per_eval=2,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bStWWyEEHHEm"
      },
      "source": [
        "## Optimize performance\n",
        "\n",
        "Techniques in TF that could achieve better performance can generally be used in TFF, e.g., [mixed precision training ](https://www.tensorflow.org/guide/mixed_precision) and [XLA](https://www.tensorflow.org/xla). The speedup (on GPUs like V100) and memory saving of mixed precision can often be significant, which could be examined by `tf.profiler`. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5p6p7hu-IEpd",
        "outputId": "1904ddfb-d699-4052-80a5-caf0405cef94"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model: \"cnn-26-4\"\n",
            "_________________________________________________________________\n",
            "Layer (type)                 Output Shape              Param #   \n",
            "=================================================================\n",
            "input_1 (InputLayer)         [(None, 28, 28, 1)]       0         \n",
            "_________________________________________________________________\n",
            "tf.image.per_image_standardi (None, 28, 28, 1)         0         \n",
            "_________________________________________________________________\n",
            "conv2d (Conv2D)              (None, 28, 28, 64)        576       \n",
            "_________________________________________________________________\n",
            "conv2d_1 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation (Activation)      (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_2 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_1 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_3 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_2 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_4 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_3 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_5 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_4 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_6 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_5 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_7 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_6 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_8 (Conv2D)            (None, 28, 28, 64)        36864     \n",
            "_________________________________________________________________\n",
            "activation_7 (Activation)    (None, 28, 28, 64)        0         \n",
            "_________________________________________________________________\n",
            "conv2d_9 (Conv2D)            (None, 14, 14, 128)       73728     \n",
            "_________________________________________________________________\n",
            "activation_8 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_10 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_9 (Activation)    (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_11 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_10 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_12 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_11 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_13 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_12 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_14 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_13 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_15 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_14 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_16 (Conv2D)           (None, 14, 14, 128)       147456    \n",
            "_________________________________________________________________\n",
            "activation_15 (Activation)   (None, 14, 14, 128)       0         \n",
            "_________________________________________________________________\n",
            "conv2d_17 (Conv2D)           (None, 7, 7, 256)         294912    \n",
            "_________________________________________________________________\n",
            "activation_16 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_18 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_17 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_19 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_18 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_20 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_19 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_21 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_20 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_22 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_21 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_23 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_22 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "conv2d_24 (Conv2D)           (None, 7, 7, 256)         589824    \n",
            "_________________________________________________________________\n",
            "activation_23 (Activation)   (None, 7, 7, 256)         0         \n",
            "_________________________________________________________________\n",
            "global_average_pooling2d (Gl (None, 256)               0         \n",
            "_________________________________________________________________\n",
            "dense (Dense)                (None, 10)                2570      \n",
            "=================================================================\n",
            "Total params: 5,827,658\n",
            "Trainable params: 5,827,658\n",
            "Non-trainable params: 0\n",
            "_________________________________________________________________\n",
            "Round 0 training loss: 2.4187185764312744, time: 18.763780117034912 secs\n",
            "Round 0 validation accuracy: 9.977468490600586\n",
            "Round 1 training loss: 2.305102825164795, time: 13.712820529937744 secs\n",
            "Round 2 training loss: 2.304737091064453, time: 9.993690172831217 secs\n",
            "Round 2 validation accuracy: 11.779976844787598\n",
            "Round 3 training loss: 2.2996833324432373, time: 9.29404467344284 secs\n",
            "Round 4 training loss: 2.299349308013916, time: 9.195427560806275 secs\n",
            "Round 4 validation accuracy: 11.224040031433105\n"
          ]
        }
      ],
      "source": [
        "# Mixed precision training. \n",
        "tff.backends.native.set_sync_local_cpp_execution_context()\n",
        "policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')\n",
        "tf.keras.mixed_precision.experimental.set_policy(policy)\n",
        "\n",
        "run_federated_training(\n",
        "    client_epochs_per_round=1, \n",
        "    train_batch_size=16, \n",
        "    test_batch_size=128, \n",
        "    cnn_num_blocks=4, \n",
        "    conv_width_multiplier=4,\n",
        "    server_learning_rate=1.0, \n",
        "    client_learning_rate=0.01,\n",
        "    total_rounds=5,\n",
        "    clients_per_round=16, \n",
        "    rounds_per_eval=2,\n",
        "    logdir='mixed'\n",
        "    )"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eiPxHOxVJAW-"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "%tensorboard --logdir=mixed --port=0"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "simulations_with_accelerators.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}