federated-learning/b - Customizing Federated Computations - Solution.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customizing Federated Computations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this next part of the tutorial it will be up to you to customize the `BaseModelOwner` or the `BaseDataOwner` to implement a new way of computing the gradients and securely aggregating them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boilerplate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First up is the boilerplate code from the previous part. This includes configuring TF Encrypted and importing all of the dependencies. We've removed the `default_model_fn` and `secure_mean` functions as you'll write new version of those."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Falling back to insecure randomness since the required custom op could not be found for the installed version of TensorFlow. Fix this by compiling custom ops. Missing file was '/Users/justinpatriquin/projects/tf-encrypted/tf_encrypted/operations/secure_random/secure_random_module_tf_2.0.0.so'\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tf_encrypted as tfe\n",
    "\n",
    "players = [\n",
    "    'server0', \n",
    "    'server1', \n",
    "    'crypto-producer', \n",
    "    'model-owner',\n",
    "    'data-owner-0',\n",
    "    'data-owner-1',\n",
    "    'data-owner-2',\n",
    "]\n",
    "config = tfe.EagerLocalConfig(players)\n",
    "\n",
    "tfe.set_config(config)\n",
    "tfe.set_protocol(tfe.protocol.Pond())\n",
    "\n",
    "from players import BaseModelOwner, BaseDataOwner\n",
    "from func_lib import default_model_fn, secure_mean, evaluate_classifier\n",
    "from util import split_dataset\n",
    "from download import download_mnist\n",
    "\n",
    "NUM_DATA_OWNERS = 3\n",
    "BATCH_SIZE = 256\n",
    "DATA_ITEMS = 60000\n",
    "BATCHES = DATA_ITEMS // NUM_DATA_OWNERS // BATCH_SIZE\n",
    "LEARNING_RATE = 0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing Reptile Meta-Learning Algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section you will use the information from the previous tutorial to help implement new functions for the `model_fn` and the `aggregator_fn`. We also recommend checking out the implementations of `default_model_fn` and `secure_mean` located in [func_lib.py](./func_lib.py) for some help figuring out where to start. \n",
    "\n",
    "We've done this with reptile and recommend following through with this but if you have another idea feel free to implement that!\n",
    "\n",
    "TODO: might need to add details to the below paragraph\n",
    "\n",
    "The reptile meta-learning algorithm computes k steps of SGD. When paired with the secure_aggregation aggregator_fn, this model_fn corresponds to using g_k as the outer gradient update. See the Reptile paper for more: https://arxiv.org/abs/1803.02999"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reptile_model_fn(data_owner, iterations=3,\n",
    "                     grad_fn=default_model_fn, **kwargs):\n",
    "  for _ in range(iterations):\n",
    "    grads_k = grad_fn(data_owner, **kwargs)\n",
    "    data_owner.optimizer.apply_gradients(\n",
    "        zip(grads_k, data_owner.model.trainable_variables),\n",
    "    )\n",
    "  return [var.read_value() for var in data_owner.model.trainable_variables]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def secure_reptile(collected_inputs, model):\n",
    "  aggr_weights = secure_mean(collected_inputs)\n",
    "\n",
    "  weights_deltas = [\n",
    "      weight - update for (weight, update) in zip(\n",
    "          model.trainable_variables, aggr_weights,\n",
    "      )\n",
    "  ]\n",
    "\n",
    "  return weights_deltas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Customize Base Classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelOwner(BaseModelOwner):\n",
    "  @classmethod\n",
    "  def model_fn(cls, data_owner):\n",
    "    return reptile_model_fn(data_owner)\n",
    "\n",
    "  @classmethod\n",
    "  def aggregator_fn(cls, model_gradients, model):\n",
    "    return secure_reptile(model_gradients, model)\n",
    "\n",
    "  @classmethod\n",
    "  def evaluator_fn(cls, model_owner):\n",
    "    return evaluate_classifier(model_owner)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO its not super clear how DataOwner should come into the picture here when customizing ModelOwner is sufficient\n",
    "class DataOwner(BaseDataOwner):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continue Boilerplate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section we continue the fill in some of the boilerplate code from the previous tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING: Splitting dataset for 3 data owners. This is for simulation use only\n"
     ]
    }
   ],
   "source": [
    "download_mnist()\n",
    "split_dataset(\"./data\", NUM_DATA_OWNERS, DATA_ITEMS)\n",
    "\n",
    "model = tf.keras.Sequential((\n",
    "    tf.keras.layers.Dense(512, input_shape=[None, 28 * 28],\n",
    "                          activation='relu'),\n",
    "    tf.keras.layers.Dense(10),\n",
    "))\n",
    "\n",
    "model.build()\n",
    "\n",
    "loss = tf.keras.losses.sparse_categorical_crossentropy\n",
    "opt = tf.keras.optimizers.Adam(LEARNING_RATE)\n",
    "\n",
    "model_owner = ModelOwner(\"model-owner\",\n",
    "                         \"{}/train.tfrecord\".format(\"./data\"),\n",
    "                         model, loss,\n",
    "                         optimizer=opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this next part consider how you might use another learning rate to customize the reptile training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simplify this with a loop?\n",
    "data_owners = [DataOwner(\"data-owner-{}\".format(i),\n",
    "                         \"{}/train{}.tfrecord\".format(\"./data\", i),\n",
    "                         model, loss,\n",
    "                         optimizer=opt)\n",
    "               for i in range(NUM_DATA_OWNERS)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now train! Remember we're using TensorFlow 2.0 so it should be  easy to explore the computations and see the actual values being passed around in the computations. You can use this to help debug any problems run into while implementing the reptile meta-learning algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "69/78 [=========================>....] - ETA: 2s - Loss: 0.7163\n",
      "Done training!!\n"
     ]
    }
   ],
   "source": [
    "model_owner.fit(data_owners, rounds=BATCHES, evaluate_every=10)\n",
    "\n",
    "print(\"\\nDone training!!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}