fedjax/experimental/notebooks/emnist_compression.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dXAJBz_5kbJW"
      },
      "source": [
        "\n",
        "[Open In Colab](https://colab.research.google.com/github/google/fedjax/blob/main/experimental/notebooks/emnist_compression.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "executionInfo": {
          "elapsed": 2,
          "status": "ok",
          "timestamp": 1626714377599,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 240
        },
        "id": "flXheaIFmfPS"
      },
      "outputs": [],
      "source": [
        "# Uncomment these to install fedjax.\n",
        "# !pip install fedjax\n",
        "# !pip install --upgrade git+https://github.com/google/fedjax.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "executionInfo": {
          "elapsed": 4901,
          "status": "ok",
          "timestamp": 1626714382615,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 240
        },
        "id": "iRWNysIeLUEX"
      },
      "outputs": [],
      "source": [
        "import fedjax\n",
        "import jax\n",
        "from jax import jit\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import optax\n",
        "import time\n",
        "from typing import Any, NamedTuple\n",
        "\n",
        "fedjax.training.set_tf_cpu_only()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyqYuFksOTnm"
      },
      "source": [
        "```\n",
        "# To disable jit, use 'with jax.disable_jit():'\n",
        "# For example,\n",
        "@jax.jit\n",
        "def f(x):\n",
        "  print(x)\n",
        "  return jnp.sum(x)\n",
        "\n",
        "x = jnp.ones([])\n",
        "\n",
        "print('jit enabled')\n",
        "for _ in range(10):\n",
        "  f(x)\n",
        "\n",
        "print('jit disabled')\n",
        "with jax.disable_jit():\n",
        "  for _ in range(10):\n",
        "    f(x)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "height": 332
        },
        "executionInfo": {
          "elapsed": 463,
          "status": "ok",
          "timestamp": 1626714383283,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 240
        },
        "id": "kp-iACjELyob",
        "outputId": "7137377a-8808-4b0d-d965-f9ea0333276c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Reusing cached file '/tmp/emnist_data/federated_emnist_train.sqlite'\n",
            "Reusing cached file '/tmp/emnist_data/federated_emnist_test.sqlite'\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "\u003cmatplotlib.image.AxesImage at 0x7f3e2ed2da10\u003e"
            ]
          },
          "execution_count": 3,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        },
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAEICAYAAACZA4KlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90\nbGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsT\nAAALEwEAmpwYAAAMd0lEQVR4nO3df4jk9X3H8efr9DzJxcJdrXIYqTax0LSgKRsbahJspMH4jwnF\non/ISYVLaYSk5I9KQtFCC7Y0CSmEwFlFm6amoYnoH5J6XFKsUKyrGH/06o+ITU4Pr8GmMaGed+e7\nf+xc2eju7Dq/6/v5gGVmv9/ZnTdz+7yZne/MflJVSHrr2zLvASTNhrFLTRi71ISxS00Yu9SEsUtN\nGLvUhLE3k6SS/DTJn03o+307yStJ7p/E99P0GHtP51fVZ098kmRvkieTvJbkmtUXzIo/TfJ8kv9O\n8k9JfvXE/qr6EPD7sxtdozJ2AXwX+APg4TX2XQH8HvABYCfwL8BXZjeaJsXYRVV9qar2A6+ssftc\n4P6qeraqjgN/C7x7pgNqIoxdG/ka8K4kv5xkK7Ab+NacZ9IITp73AFp4h4B/Bp4EjgM/AD4014k0\nEu/ZtZEbgPcCZwOnAn8CfDvJ2+Y6ld40Y9dGzgf+vqoOVtWxqroN2IG/t/+/Y+wiySlJTgUCbE1y\napITPxsPAlckOTPJliRXA1uBZ+Y1r0Zj7AK4F/gf4DeBvYPzHxzs+3NWDs09AvwI+EPgd6rqR7Me\nUuOJf6mmlySvAEeAv6qqP57A99sHvA/416q6ZNzvp+kxdqkJH8ZLTRi71MRMX1RzSrbVqWyf5VVK\nrbzCT3m1jmStfWPFnuRS4IvAScBfV9VNwy5/Ktv5jfgcjjQtD9T+dfeN/DA+yUnAl4CPsPICi6uS\n+EILaUGN8zv7hcAzg3dDvcrKGyYun8xYkiZtnNjPYuVNESccHGz7GUn2JFlOsnyUI2NcnaRxjBP7\nWk8CvOGgfVXtraqlqlrayrYxrk7SOMaJ/SAr74Q64R3AC+ONI2laxon9QeC8JOcmOQW4Erh7MmNJ\nmrSRD71V1bEk1wH/yMqht1ur6omJTSZposY6zl5V9wD3TGgWSVPky2WlJoxdasLYpSaMXWrC2KUm\njF1qwtilJoxdasLYpSaMXWrC2KUmjF1qwtilJoxdasLYpSaMXWrC2KUmjF1qwtilJoxdasLYpSaM\nXWrC2KUmjF1qwtilJoxdasLYpSaMXWrC2KUmjF1qYqwlm5M8B7wMHAeOVdXSJIaSNHljxT7wW1X1\nwwl8H0lT5MN4qYlxYy/g3iQPJdmz1gWS7EmynGT5KEfGvDpJoxr3YfxFVfVCkjOAfUn+varuW32B\nqtoL7AX4ueysMa9P0ojGumevqhcGp4eBO4ELJzGUpMkbOfYk25OcduI88GHg8UkNJmmyxnkYfyZw\nZ5IT3+fvqupbE5lK0sSNHHtVPQucP8FZJE2Rh96kJoxdasLYpSaMXWrC2KUmjF1qwtilJoxdasLY\npSaMXWrC2KUmjF1qwtilJoxdasLYpSaMXWrC2KUmjF1qwtilJoxdasLYpSaMXWrC2KUmjF1qwtil\nJoxdasLYpSaMXWrC2KUmxlmyWR1sOWn4/teOz2YOjW3De/YktyY5nOTxVdt2JtmX5OnB6Y7pjilp\nXJt5GH8bcOnrtl0P7K+q84D9g88lLbANY6+q+4CXXrf5cuD2wfnbgY9OdixJkzbqE3RnVtUhgMHp\nGetdMMmeJMtJlo9yZMSrkzSuqT8bX1V7q2qpqpa2sm3aVydpHaPG/mKSXQCD08OTG0nSNIwa+93A\n7sH53cBdkxlH0rRseJw9yR3AxcDpSQ4CNwA3AV9Pci3wfeCKaQ6pOfI4+lvGhrFX1VXr7LpkwrNI\nmiJfLis1YexSE8YuNWHsUhPGLjXhW1zfArL1lHX31dFXh37tUze/d4Pv/drQ/edd89DQ/UPfIuth\nvZnynl1qwtilJoxdasLYpSaMXWrC2KUmjF1qwuPszW15efifil6+4otD91+55QPDr8Bj6QvDe3ap\nCWOXmjB2qQljl5owdqkJY5eaMHapCY+zN7fjQIbu35bhPyInveucofuPP/W99Xe6HPRMec8uNWHs\nUhPGLjVh7FITxi41YexSE8YuNeFx9reCGv633Yc59b+Gf+3btqz/N+kBOHmDY+VaGBvesye5Ncnh\nJI+v2nZjkueTPDL4uGy6Y0oa12Yext8GXLrG9i9U1QWDj3smO5akSdsw9qq6D3hpBrNImqJxnqC7\nLsmjg4f5O9a7UJI9SZaTLB/lyBhXJ2kco8b+ZeCdwAXAIeBz612wqvZW1VJVLW1l24hXJ2lcI8Ve\nVS9W1fGqeg24GbhwsmNJmrSRYk+ya9WnHwMeX++ykhbDhsfZk9wBXAycnuQgcANwcZILgAKeAz4+\nvRG1kTp2bOSv3f4PDwzd//4N/mlPe+6xka/b96vP1oaxV9VVa2y+ZQqzSJoiXy4rNWHsUhPGLjVh\n7FITxi414Vtcu8vwPyW90aG50d9cq1nznl1qwtilJoxdasLYpSaMXWrC2KUmjF1qwuPs3VUN3Z2T\nh/+IjPP2Ws2W9+xSE8YuNWHsUhPGLjVh7FITxi41YexSEx5n11AeR3/r8J5dasLYpSaMXWrC2KUm\njF1qwtilJoxdamLD2JOcneQ7SQ4keSLJJwfbdybZl+TpwemO6Y8raVSbuWc/Bny6qn4FeB/wiSTv\nBq4H9lfVecD+weeSFtSGsVfVoap6eHD+ZeAAcBZwOXD74GK3Ax+d0oySJuBN/c6e5BzgPcADwJlV\ndQhW/kMAzpj4dJImZtOxJ3k78A3gU1X14zfxdXuSLCdZPsqRUWaUNAGbij3JVlZC/2pVfXOw+cUk\nuwb7dwGH1/raqtpbVUtVtbSVbZOYWdIINvNsfIBbgANV9flVu+4Gdg/O7wbumvx4kiZlM29xvQi4\nGngsySODbZ8BbgK+nuRa4PvAFVOZUNJEbBh7Vd0PrLeI9yWTHUfStPgKOqkJY5eaMHapCWOXmjB2\nqQljl5owdqkJY5eaMHapCWOXmjB2qQljl5owdqkJY5eaMHapCWOXmjB2qQljl5owdqkJY5eaMHap\nCWOXmjB2qQljl5owdqkJY5eaMHapCWOXmjB2qQljl5rYMPYkZyf5TpIDSZ5I8snB9huTPJ/kkcHH\nZdMfV9KoNlyfHTgGfLqqHk5yGvBQkn2DfV+oqr+c3niSJmXD2KvqEHBocP7lJAeAs6Y9mKTJelO/\nsyc5B3gP8MBg03VJHk1ya5Id63zNniTLSZaPcmS8aSWNbNOxJ3k78A3gU1X1Y+DLwDuBC1i55//c\nWl9XVXuraqmqlraybfyJJY1kU7En2cpK6F+tqm8CVNWLVXW8ql4DbgYunN6Yksa1mWfjA9wCHKiq\nz6/avmvVxT4GPD758SRNymaejb8IuBp4LMkjg22fAa5KcgFQwHPAx6cwn6QJ2cyz8fcDWWPXPZMf\nR9K0+Ao6qQljl5owdqkJY5eaMHapCWOXmjB2qQljl5owdqkJY5eaMHapCWOXmjB2qQljl5pIVc3u\nypL/BP5j1abTgR/ObIA3Z1FnW9S5wNlGNcnZfrGqfmGtHTON/Q1XnixX1dLcBhhiUWdb1LnA2UY1\nq9l8GC81YexSE/OOfe+cr3+YRZ1tUecCZxvVTGab6+/skmZn3vfskmbE2KUm5hJ7kkuTPJnkmSTX\nz2OG9SR5Lsljg2Wol+c8y61JDid5fNW2nUn2JXl6cLrmGntzmm0hlvEessz4XG+7eS9/PvPf2ZOc\nBDwF/DZwEHgQuKqq/m2mg6wjyXPAUlXN/QUYST4I/AT4m6r6tcG2vwBeqqqbBv9R7qiqP1qQ2W4E\nfjLvZbwHqxXtWr3MOPBR4BrmeNsNmet3mcHtNo979guBZ6rq2ap6FfgacPkc5lh4VXUf8NLrNl8O\n3D44fzsrPywzt85sC6GqDlXVw4PzLwMnlhmf6203ZK6ZmEfsZwE/WPX5QRZrvfcC7k3yUJI98x5m\nDWdW1SFY+eEBzpjzPK+34TLes/S6ZcYX5rYbZfnzcc0j9rWWklqk438XVdWvAx8BPjF4uKrN2dQy\n3rOyxjLjC2HU5c/HNY/YDwJnr/r8HcALc5hjTVX1wuD0MHAni7cU9YsnVtAdnB6e8zz/Z5GW8V5r\nmXEW4Lab5/Ln84j9QeC8JOcmOQW4Erh7DnO8QZLtgydOSLId+DCLtxT13cDuwfndwF1znOVnLMoy\n3ustM86cb7u5L39eVTP/AC5j5Rn57wGfnccM68z1S8B3Bx9PzHs24A5WHtYdZeUR0bXAzwP7gacH\npzsXaLavAI8Bj7IS1q45zfZ+Vn41fBR4ZPBx2bxvuyFzzeR28+WyUhO+gk5qwtilJoxdasLYpSaM\nXWrC2KUmjF1q4n8BpRbJRivOFnQAAAAASUVORK5CYII=\n",
            "text/plain": [
              "\u003cFigure size 600x400 with 1 Axes\u003e"
            ]
          },
          "metadata": {
            "tags": []
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "# Load the model and the data.\n",
        "model = fedjax.models.emnist.create_conv_model(only_digits=False)\n",
        "train, test = fedjax.datasets.emnist.load_data(only_digits=False,\n",
        "                                               cache_dir='/tmp/emnist_data')\n",
        "# Print a single client data to verify data is loaded.\n",
        "example_client_id = list(train.client_ids())[0]\n",
        "single_client_data = list(train.get_client(\n",
        "    example_client_id).batch(batch_size=1))\n",
        "pixels = single_client_data[0]['x']\n",
        "label = single_client_data[0]['y']\n",
        "plt.title(f'{label}')\n",
        "plt.imshow(pixels.reshape(28, 28))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "executionInfo": {
          "elapsed": 66,
          "status": "ok",
          "timestamp": 1626714383620,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 240
        },
        "id": "8s30qnrML6t9"
      },
      "outputs": [],
      "source": [
        "@fedjax.dataclass\n",
        "class ServerState:\n",
        "  params : fedjax.Params\n",
        "  opt_state : fedjax.optimizers.OptState\n",
        "  aggregator_state: Any\n",
        "\n",
        "# Define federated algorithm.\n",
        "def fed_avg(grad_fn,\n",
        "            client_optimizer,\n",
        "            server_optimizer,\n",
        "            client_batch_hparams,\n",
        "            aggregator):\n",
        "  \"\"\"Builds the basic implementation of federated averaging.\"\"\"\n",
        "\n",
        "  def init(params):\n",
        "    \"\"\" Initialize the federated algorithm\"\"\"\n",
        "    opt_state = server_optimizer.init(params)\n",
        "    aggregator_state = aggregator.init()\n",
        "    return ServerState(params, opt_state, aggregator_state)\n",
        "\n",
        "  def apply(server_state, clients): \n",
        "    \"\"\" Apply one round of federated algorithm/\"\"\"\n",
        "    client_delta_params_weights = []\n",
        "    for client_id, client_dataset, client_rng in clients:\n",
        "      delta_params, client_loss = client_update(\n",
        "          server_state.params, client_dataset, client_rng)\n",
        "      client_delta_params_weights.append((client_id, \n",
        "                                          delta_params,\n",
        "                                          len(client_dataset)))\n",
        "    weighted_averaged_delta, new_aggregator_state = aggregator.apply(\n",
        "        client_delta_params_weights,\n",
        "        server_state.aggregator_state)\n",
        "    server_state = server_update(server_state,\n",
        "                                 weighted_averaged_delta,\n",
        "                                 new_aggregator_state)\n",
        "    return server_state\n",
        "\n",
        "  def client_update(server_params, client_dataset, client_rng):\n",
        "    params = server_params\n",
        "    opt_state = client_optimizer.init(params)\n",
        "    padded_batches = fedjax.padded_batch_client_datasets(\n",
        "        [client_dataset], batch_size=1024)\n",
        "    metrics = fedjax.evaluate_model(model, params, padded_batches)\n",
        "    client_loss = metrics['loss']\n",
        "    for batch in client_dataset.shuffle_repeat_batch(client_batch_hparams):\n",
        "      client_rng, use_rng = jax.random.split(client_rng)\n",
        "      grads = grad_fn(params, batch, use_rng)\n",
        "      opt_state, params = client_optimizer.apply(grads, opt_state, params)\n",
        "    delta_params = jax.tree_util.tree_multimap(lambda a, b: a - b,\n",
        "                                               server_params, params)\n",
        "    return delta_params, client_loss\n",
        "\n",
        "  def server_update(server_state, mean_delta_params, new_aggregator_state):\n",
        "    opt_state, params = server_optimizer.apply(mean_delta_params,\n",
        "                                               server_state.opt_state,\n",
        "                                               server_state.params)\n",
        "    return ServerState(params, opt_state, new_aggregator_state)\n",
        "\n",
        "  return fedjax.FederatedAlgorithm(init, apply)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "executionInfo": {
          "elapsed": 527294,
          "status": "ok",
          "timestamp": 1626714911066,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 240
        },
        "id": "O5IHdAjEjOan",
        "outputId": "e508884d-4ee6-4a52-e692-bf2a6afde001"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "round_num: 100, mean_round_duration: 0.4552804470062256 sec\n",
            "round_num: 200, mean_round_duration: 0.3484324908256531 sec\n",
            "round_num: 300, mean_round_duration: 0.3398168110847473 sec\n",
            "round_num: 400, mean_round_duration: 0.34802984476089477 sec\n",
            "round_num: 500, mean_round_duration: 0.3342078638076782 sec\n",
            "round_num: 600, mean_round_duration: 0.33866923570632934 sec\n",
            "round_num: 700, mean_round_duration: 0.34813843011856077 sec\n",
            "round_num: 800, mean_round_duration: 0.3442531204223633 sec\n",
            "round_num: 900, mean_round_duration: 0.3395417761802673 sec\n",
            "round_num: 1000, mean_round_duration: 0.338899233341217 sec\n",
            "round_num: 1100, mean_round_duration: 0.33501195192337035 sec\n",
            "round_num: 1200, mean_round_duration: 0.34103004932403563 sec\n",
            "round_num: 1300, mean_round_duration: 0.338041832447052 sec\n",
            "round_num: 1400, mean_round_duration: 0.3409284591674805 sec\n",
            "round_num: 1500, mean_round_duration: 0.33965426445007324 sec\n",
            "{'accuracy': DeviceArray(0.8554393, dtype=float32), 'loss': DeviceArray(0.42102018, dtype=float32)}\n"
          ]
        }
      ],
      "source": [
        "# Run federated algorithm.\n",
        "# The hyper-parameters here is not the state of the art. It achieves a\n",
        "# global test accuracy as follows\n",
        "# num_levels, server_learning_rate_starting_point, test_accuracy \n",
        "# 256, 1.0, 85.6%\n",
        "# 2, 0.01, 59.2%\n",
        "\n",
        "num_rounds = 1500\n",
        "num_clients_per_round = 10\n",
        "client_batch_size = 10\n",
        "client_optimizer = fedjax.optimizers.sgd(learning_rate=0.16)\n",
        "server_schedule = optax.exponential_decay(1.0, 500, 0.5)\n",
        "server_optimizer = fedjax.optimizers.sgd(server_schedule)\n",
        "client_batch_hparams = fedjax.ShuffleRepeatBatchHParams(\n",
        "    batch_size=10,num_epochs=1)\n",
        "grad_fn = fedjax.model_grad(model)\n",
        "# Creating a federated algorithm with 8 bits of quantization.\n",
        "algorithm = fed_avg(grad_fn, \n",
        "                    client_optimizer,\n",
        "                    server_optimizer,\n",
        "                    client_batch_hparams,\n",
        "                    fedjax.aggregators.uniform_stochastic_quantizer(\n",
        "                        256, jax.random.PRNGKey(0)))\n",
        "init_params = model.init(jax.random.PRNGKey(0))\n",
        "rng = jax.random.PRNGKey(42)\n",
        "server_state = algorithm.init(init_params)\n",
        "sampler = fedjax.client_samplers.UniformGetClientSampler(\n",
        "    train, num_clients_per_round, seed=0)\n",
        "start = time.time()\n",
        "for round_num in range(1, num_rounds + 1):\n",
        "  clients = sampler.sample()\n",
        "  server_state = algorithm.apply(server_state, clients)\n",
        "  if round_num % 100 == 0:\n",
        "    print('round_num: {}, mean_round_duration: {} sec'.format(\n",
        "        round_num, (time.time() - start)/100))\n",
        "    start = time.time() \n",
        "final_params = server_state.params\n",
        "\n",
        "overall_metrics = fedjax.evaluate_model(\n",
        "    model,\n",
        "    final_params,\n",
        "    fedjax.padded_batch_federated_data(test, batch_size=1024))\n",
        "print(overall_metrics)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      },
      "name": "emnist-compression.ipynb",
      "provenance": [
        {
          "file_id": "1d3YXAA7t9NFgWFp6YyUNiDHRAGbCaIwq",
          "timestamp": 1619130258104
        },
        {
          "file_id": "1UhglgvY6rQOPGQQPIL7aJRQ5wXndCrWt",
          "timestamp": 1617290897629
        },
        {
          "file_id": "1p7KiuZnXC3KOeJW_ArP-cTzpax0IBBFZ",
          "timestamp": 1616023981580
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}