Part_4_Encrypted_Federated_Learning_using_PySyft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# [Part 4] - Encrypted Federated Learning using PySyft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we're using the structure of the dataset [A Pressure Map Data set for In-bed Posture Classification](https://physionet.org/content/pmd/1.0.0/) to simulate a situation in which each subject has collected their data with the use of a mobile app. Our intention is to train a global model making use of each individual's data, without having the need to see or move the data out from their devices. For this, we're going to use Federated Learning, which allows us to train individual models inside each device, and average them together to create a global model. This way, no data is revealed to any third party."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have a folder called __dataset__, which contains two experiments (downloaded in part 1). In these project we're only using 'Experiment I'."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0820 13:26:00.819633 4394591680 secure_random.py:26] 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 '/anaconda3/envs/pysyft/lib/python3.7/site-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.14.0.so'\n",
      "W0820 13:26:00.834891 4394591680 deprecation_wrapper.py:119] From /anaconda3/envs/pysyft/lib/python3.7/site-packages/tf_encrypted/session.py:26: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from itertools import islice\n",
    "import torch as th\n",
    "import syft as sy\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Ar0HQVg344ZO"
   },
   "source": [
    "## Building our subjects database (revisited)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hPRIons74-oP"
   },
   "source": [
    "We want to simulate a situation where all data is stored inside each subject's device. For this, we're first going to retrieve the data and store it in a dictionary. The structure will be `{\"SubjectName\": (dataTensor, labelsTensor)}`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aGvpj6Ft44AB"
   },
   "outputs": [],
   "source": [
    "import torch as th\n",
    "\n",
    "path_exp_1 = \"dataset/experiment-i\"\n",
    "positions = [\"justAPlaceholder\",\"supine\", \"right\", \"left\", \"right\", \"right\", \"left\", \"left\", \"supine\", \"supine\", \"supine\", \"supine\", \"supine\", \"right\", \"left\", \"supine\", \"supine\", \"supine\"]\n",
    "\n",
    "subjects_dict = {}\n",
    "\n",
    "for _, dirs, _ in os.walk(path_exp_1):\n",
    "  for directory in dirs:\n",
    "    # each directory is a subject\n",
    "    subject = directory\n",
    "    data = None\n",
    "    labels = None\n",
    "    \n",
    "    for _, _, files in os.walk(os.path.join(path_exp_1, directory)):\n",
    "      for file in files:\n",
    "        file_path = os.path.join(path_exp_1, directory, file)\n",
    "        with open(file_path, 'r') as f:\n",
    "          # Start from second recording, as the first two are corrupted\n",
    "          for line in f.read().splitlines()[2:]:\n",
    "            def token_position(x):\n",
    "              return {\n",
    "                'supine': 0,\n",
    "                'left': 1,\n",
    "                'right': 2,\n",
    "                'left_fetus': 1,\n",
    "                'right_fetus': 2\n",
    "              }[x]\n",
    "            \n",
    "            \n",
    "            raw_data = np.fromstring(line, dtype=float, sep='\\t').reshape(1,64,32)\n",
    "            file_data = np.round(raw_data*255/1000).astype(np.uint8) # Change the range from [0-1000] to [0-255]. This allows us to use tranforms later.\n",
    "            file_label = token_position(positions[int(file[:-4])]) # Turn the file index into position list, and turn position list into reduced indices.\n",
    "            file_label = np.array([file_label])\n",
    "            \n",
    "            if data is None:\n",
    "              data = file_data\n",
    "            else:\n",
    "              data = np.concatenate((data, file_data), axis=0)\n",
    "\n",
    "            if labels is None:\n",
    "              labels = file_label\n",
    "            else:\n",
    "              labels = np.concatenate((labels, file_label), axis=0)\n",
    "              \n",
    "    subjects_dict[subject] = (th.from_numpy(data), th.from_numpy(labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "z4HJS9K-G8RJ"
   },
   "source": [
    "## Create the workers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UzswwrE1G-Dq"
   },
   "source": [
    "Now that we have our data partitioned, let's create a worker for each subject. We'll also create a secure worker, which will provide encryption mechanisms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "slMvzqhmIeYl"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "# We don't want the download process to fill our screens\n",
    "!pip install syft\n",
    "import syft as sy\n",
    "\n",
    "hook = sy.TorchHook(th)\n",
    "\n",
    "workers = [sy.VirtualWorker(id=key, hook=hook) for key in subjects_dict.keys()]\n",
    "\n",
    "secure_worker = sy.VirtualWorker(id=\"secure_worker\", hook=hook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vQs1wPFYZjXo"
   },
   "source": [
    "## Create the datasets and dataloaders"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Rgfshk0wZnfd"
   },
   "source": [
    "Let's create a dataset and their respective trainloader and testloader for each subject"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "7IiET9P6ZgFm"
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "# Let's create a copy of the data dict. This way the original dict is untouched.\n",
    "subj_dict = copy.deepcopy(subjects_dict)\n",
    "datasets = {}\n",
    "trainloaders = {}\n",
    "testloaders = {}\n",
    "train_percent = 0.8\n",
    "\n",
    "for worker in workers:\n",
    "  # Create the Dataset\n",
    "  datasets[worker.id] = sy.BaseDataset(*subj_dict[worker.id])\n",
    "  \n",
    "  train_size = int(train_percent * len(datasets[worker.id]))\n",
    "  test_size = len(datasets[worker.id]) - train_size\n",
    "  \n",
    "  # Split the dataset for the dataloaders\n",
    "  train_dataset, test_dataset = th.utils.data.random_split(datasets[worker.id], [train_size, test_size])\n",
    "  \n",
    "  trainloaders[worker.id] = th.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
    "  testloaders[worker.id] = th.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
    "  \n",
    "  # Send the dataset to the worker\n",
    "  datasets[worker.id] = datasets[worker.id].send(worker)\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "nxAsWFPguYNb"
   },
   "source": [
    "## Define the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "meCZfEj9ubIG"
   },
   "source": [
    "This model will be copied and sent to every subject, trained locally, and then we obtain the average through the use of Additive Secret Sharing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 187
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 315,
     "status": "ok",
     "timestamp": 1565377009424,
     "user": {
      "displayName": "Alejandro Aristizabal",
      "photoUrl": "https://lh6.googleusercontent.com/-nfBO6mFL4Mk/AAAAAAAAAAI/AAAAAAAAL1o/9ct9fKu4tjk/s64/photo.jpg",
      "userId": "17089498558649987582"
     },
     "user_tz": 300
    },
    "id": "C2vMiFmAJkoc",
    "outputId": "63613f91-d3f6-4c7a-8fef-0b003d6063e4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Network(\n",
       "  (conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (conv2): Conv2d(6, 18, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  (h1): Linear(in_features=2304, out_features=392, bias=True)\n",
       "  (h2): Linear(in_features=392, out_features=98, bias=True)\n",
       "  (h3): Linear(in_features=98, out_features=3, bias=True)\n",
       "  (relu): ReLU()\n",
       "  (logsoftmax): LogSoftmax()\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class Network(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        #Input channels = 1, output channels = 6\n",
    "        self.conv1 = torch.nn.Conv2d(1, 6, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv2 = torch.nn.Conv2d(6, 18, kernel_size=3, stride=1, padding=1)\n",
    "        \n",
    "        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0)\n",
    "        \n",
    "        \n",
    "        # Inputs to hidden layer linear transformation\n",
    "        self.h1 = nn.Linear(18 * 16 * 8, 392)\n",
    "        self.h2 = nn.Linear(392, 98)\n",
    "        self.h3 = nn.Linear(98, 3)\n",
    "        \n",
    "        # Define sigmoid activation and softmax output \n",
    "        self.relu = nn.ReLU()\n",
    "        self.logsoftmax = nn.LogSoftmax(dim=1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = x.float()\n",
    "        # Add a \"channel dimension\"\n",
    "        x = x.unsqueeze(1)\n",
    "        \n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = self.pool(x)\n",
    "        x = F.relu(self.conv2(x))\n",
    "        x = self.pool(x)\n",
    "        \n",
    "        x = x.view(x.shape[0], -1)\n",
    "\n",
    "        x = self.relu(self.h1(x))\n",
    "        x = self.relu(self.h2(x))\n",
    "        x = self.logsoftmax(self.h3(x))\n",
    "        \n",
    "        return x\n",
    "      \n",
    "net = Network()\n",
    "net"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L_LoEReVuzp0"
   },
   "source": [
    "## Train on each subject"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a model and each subject containing their private datasets. For training we're going to make a copy of the model to each individual, and train remotely using their dataset. It's important to remember that since each dataset is composed of a single individual's examples, it's highly probable that the training will result in overfitting. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MwCPYOpTQmhH"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device = \"cpu\" # TODO: research if cuda is supported on federated learning\n",
    "net.to(device)\n",
    "\n",
    "models = {}\n",
    "\n",
    "for worker in workers:\n",
    "  models[worker.id] = net.copy().send(worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 680
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 94489,
     "status": "ok",
     "timestamp": 1565377106083,
     "user": {
      "displayName": "Alejandro Aristizabal",
      "photoUrl": "https://lh6.googleusercontent.com/-nfBO6mFL4Mk/AAAAAAAAAAI/AAAAAAAAL1o/9ct9fKu4tjk/s64/photo.jpg",
      "userId": "17089498558649987582"
     },
     "user_tz": 300
    },
    "id": "LeFk_AAHpQhy",
    "outputId": "3adcd951-bc97-4011-a98c-6ba717b6a875"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subject S5.. Epoch 1/3.. Train loss: 0.243.. Test loss: 0.028.. Test accuracy: 1.000\n",
      "Subject S5.. Epoch 2/3.. Train loss: 0.015.. Test loss: 0.008.. Test accuracy: 1.000\n",
      "Subject S5.. Epoch 3/3.. Train loss: 0.006.. Test loss: 0.004.. Test accuracy: 1.000\n",
      "Subject S2.. Epoch 1/3.. Train loss: 0.500.. Test loss: 0.150.. Test accuracy: 1.000\n",
      "Subject S2.. Epoch 2/3.. Train loss: 0.092.. Test loss: 0.049.. Test accuracy: 1.000\n",
      "Subject S2.. Epoch 3/3.. Train loss: 0.034.. Test loss: 0.023.. Test accuracy: 1.000\n",
      "Subject S3.. Epoch 1/3.. Train loss: 0.435.. Test loss: 0.077.. Test accuracy: 1.000\n",
      "Subject S3.. Epoch 2/3.. Train loss: 0.036.. Test loss: 0.016.. Test accuracy: 1.000\n",
      "Subject S3.. Epoch 3/3.. Train loss: 0.011.. Test loss: 0.007.. Test accuracy: 1.000\n",
      "Subject S4.. Epoch 1/3.. Train loss: 0.309.. Test loss: 0.057.. Test accuracy: 1.000\n",
      "Subject S4.. Epoch 2/3.. Train loss: 0.026.. Test loss: 0.014.. Test accuracy: 1.000\n",
      "Subject S4.. Epoch 3/3.. Train loss: 0.009.. Test loss: 0.007.. Test accuracy: 1.000\n",
      "Subject S10.. Epoch 1/3.. Train loss: 0.488.. Test loss: 0.099.. Test accuracy: 1.000\n",
      "Subject S10.. Epoch 2/3.. Train loss: 0.060.. Test loss: 0.024.. Test accuracy: 1.000\n",
      "Subject S10.. Epoch 3/3.. Train loss: 0.017.. Test loss: 0.010.. Test accuracy: 1.000\n",
      "Subject S11.. Epoch 1/3.. Train loss: 0.415.. Test loss: 0.097.. Test accuracy: 0.992\n",
      "Subject S11.. Epoch 2/3.. Train loss: 0.045.. Test loss: 0.030.. Test accuracy: 0.992\n",
      "Subject S11.. Epoch 3/3.. Train loss: 0.016.. Test loss: 0.019.. Test accuracy: 0.996\n",
      "Subject S8.. Epoch 1/3.. Train loss: 0.315.. Test loss: 0.053.. Test accuracy: 1.000\n",
      "Subject S8.. Epoch 2/3.. Train loss: 0.031.. Test loss: 0.014.. Test accuracy: 1.000\n",
      "Subject S8.. Epoch 3/3.. Train loss: 0.011.. Test loss: 0.007.. Test accuracy: 1.000\n",
      "Subject S1.. Epoch 1/3.. Train loss: 0.391.. Test loss: 0.064.. Test accuracy: 1.000\n",
      "Subject S1.. Epoch 2/3.. Train loss: 0.033.. Test loss: 0.016.. Test accuracy: 1.000\n",
      "Subject S1.. Epoch 3/3.. Train loss: 0.010.. Test loss: 0.008.. Test accuracy: 1.000\n",
      "Subject S6.. Epoch 1/3.. Train loss: 0.405.. Test loss: 0.106.. Test accuracy: 1.000\n",
      "Subject S6.. Epoch 2/3.. Train loss: 0.056.. Test loss: 0.025.. Test accuracy: 1.000\n",
      "Subject S6.. Epoch 3/3.. Train loss: 0.017.. Test loss: 0.012.. Test accuracy: 1.000\n",
      "Subject S7.. Epoch 1/3.. Train loss: 0.255.. Test loss: 0.025.. Test accuracy: 0.992\n",
      "Subject S7.. Epoch 2/3.. Train loss: 0.016.. Test loss: 0.013.. Test accuracy: 0.994\n",
      "Subject S7.. Epoch 3/3.. Train loss: 0.010.. Test loss: 0.011.. Test accuracy: 1.000\n",
      "Subject S9.. Epoch 1/3.. Train loss: 0.374.. Test loss: 0.073.. Test accuracy: 1.000\n",
      "Subject S9.. Epoch 2/3.. Train loss: 0.033.. Test loss: 0.013.. Test accuracy: 1.000\n",
      "Subject S9.. Epoch 3/3.. Train loss: 0.008.. Test loss: 0.006.. Test accuracy: 1.000\n",
      "Subject S13.. Epoch 1/3.. Train loss: 0.430.. Test loss: 0.140.. Test accuracy: 1.000\n",
      "Subject S13.. Epoch 2/3.. Train loss: 0.070.. Test loss: 0.033.. Test accuracy: 1.000\n",
      "Subject S13.. Epoch 3/3.. Train loss: 0.020.. Test loss: 0.014.. Test accuracy: 1.000\n",
      "Subject S12.. Epoch 1/3.. Train loss: 0.198.. Test loss: 0.017.. Test accuracy: 1.000\n",
      "Subject S12.. Epoch 2/3.. Train loss: 0.008.. Test loss: 0.004.. Test accuracy: 1.000\n",
      "Subject S12.. Epoch 3/3.. Train loss: 0.003.. Test loss: 0.002.. Test accuracy: 1.000\n"
     ]
    }
   ],
   "source": [
    "from torch import optim\n",
    "\n",
    "for worker in workers:\n",
    "  \n",
    "  model = models[worker.id]\n",
    "  trainloader = trainloaders[worker.id]\n",
    "  testloader = testloaders[worker.id]\n",
    "  \n",
    "  criterion = nn.NLLLoss()\n",
    "  optimizer = optim.Adam(model.parameters(), lr = 0.00003)\n",
    "  \n",
    "  epochs = 3\n",
    "  running_loss = 0\n",
    "  train_losses, test_losses = [], []\n",
    "  \n",
    "  for epoch in range(epochs):\n",
    "    for inputs, labels in trainloader:\n",
    "      \n",
    "      inputs, labels = inputs.to(device), labels.to(device)\n",
    "      \n",
    "      optimizer.zero_grad()\n",
    "      \n",
    "      logps = model.forward(inputs)\n",
    "      loss = criterion(logps, labels)\n",
    "      loss.backward()\n",
    "      optimizer.step()\n",
    "      \n",
    "      running_loss += loss\n",
    "      \n",
    "    else:\n",
    "      \n",
    "      test_loss = 0\n",
    "      accuracy = 0\n",
    "      model.eval()\n",
    "    \n",
    "      with torch.no_grad():\n",
    "        for inputs, labels in testloader:\n",
    "          inputs, labels = inputs.to(device), labels.to(device)\n",
    "          logps = model.forward(inputs)\n",
    "          test_loss += criterion(logps, labels)\n",
    "        \n",
    "          ps = torch.exp(logps)\n",
    "          top_p, top_class = ps.topk(1, dim=1)\n",
    "          equals = top_class == labels.view(*top_class.shape)\n",
    "          accuracy += torch.mean(equals.float())\n",
    "    \n",
    "      train_losses.append(running_loss/len(trainloader))\n",
    "      test_losses.append(test_loss/len(testloader))\n",
    "        \n",
    "      print(f\"Subject {worker.id}.. \"\n",
    "            f\"Epoch {epoch+1}/{epochs}.. \"\n",
    "            f\"Train loss: {running_loss.get()/len(trainloader):.3f}.. \"\n",
    "            f\"Test loss: {test_loss.get()/len(testloader):.3f}.. \"\n",
    "            f\"Test accuracy: {accuracy.get()/len(testloader):.3f}\")\n",
    "      running_loss = 0\n",
    "      model.train()\n",
    "      "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Anw8nZwU0f8m"
   },
   "source": [
    "# Combine the models into a global model and update local models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to take advantage of what each model learned so that we get a better, more generalized model for all our subjects. For this, we're going to average all of the models' parameters and use this average to build our global model. This average should result in a model with good accuracy for all subjects. Once we build our global model, we can send it back to each subject so that all devices take advantage of our generalized model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "21LyQct3uCvd"
   },
   "source": [
    "## Share the parameter's of each subject's model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "VMwSTB3MuMUD"
   },
   "source": [
    "We use Additive Secret Sharing to allow operations to be performed between models and tensors without sacrificing privacy. This way we can safely combine all parameters without compromising any data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "0A53P_Z-7U2Y"
   },
   "outputs": [],
   "source": [
    "# There seems to be a bug with fix_prec() on model parameters. (https://github.com/OpenMined/PySyft/issues/2490)\n",
    "# For this reason we have to do fix_prec().share() directly inside the pointers\n",
    "\n",
    "for worker in workers:\n",
    "  model = models[worker.id]\n",
    "  for p in model.parameters():\n",
    "    # This is equivalent to model.fix_precision().share(*workers, crypto_provider=secure_worker)\n",
    "    p.data = p.data.fix_precision().share(*workers, crypto_provider=secure_worker)\n",
    "    \n",
    "  models[worker.id] = model\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 272
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 304,
     "status": "ok",
     "timestamp": 1565377125328,
     "user": {
      "displayName": "Alejandro Aristizabal",
      "photoUrl": "https://lh6.googleusercontent.com/-nfBO6mFL4Mk/AAAAAAAAAAI/AAAAAAAAL1o/9ct9fKu4tjk/s64/photo.jpg",
      "userId": "17089498558649987582"
     },
     "user_tz": 300
    },
    "id": "LVUWOYUXt3fG",
    "outputId": "b5cbc014-4ed7-4d50-9c78-827eea50d414"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Wrapper)>FixedPrecisionTensor>[AdditiveSharingTensor]\n",
       "\t-> [PointerTensor | me:39500411724 -> S5:26126858823]\n",
       "\t-> [PointerTensor | me:14127617727 -> S2:47416863968]\n",
       "\t-> [PointerTensor | me:67162268495 -> S3:68678061802]\n",
       "\t-> [PointerTensor | me:88426740161 -> S4:54765261537]\n",
       "\t-> [PointerTensor | me:40428048369 -> S10:97459749673]\n",
       "\t-> [PointerTensor | me:76059908807 -> S11:68114333795]\n",
       "\t-> [PointerTensor | me:7485253650 -> S8:14829575712]\n",
       "\t-> [PointerTensor | me:33759476462 -> S1:29512608699]\n",
       "\t-> [PointerTensor | me:54708755841 -> S6:26236504876]\n",
       "\t-> [PointerTensor | me:57632845670 -> S7:80817756787]\n",
       "\t-> [PointerTensor | me:40146488240 -> S9:4586128961]\n",
       "\t-> [PointerTensor | me:40712918845 -> S13:89769970463]\n",
       "\t-> [PointerTensor | me:27077482814 -> S12:46688925286]\n",
       "\t*crypto provider: secure_worker*"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Sanity Check\n",
    "\n",
    "param = list(models['S1'].parameters())[0]\n",
    "\n",
    "# This should be (Wrapper)>FixedPrecisionTensor>[AdditiveSharingTensor]\n",
    "param.location._objects[param.id_at_location]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oAAtZf4ofRfM"
   },
   "source": [
    "## Move the encrypted models to our secure worker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "oyp55kyufizH"
   },
   "source": [
    "Note that the data is still encrypted and shared among all devices. This is only so that all the aggregation is executed by the secure_worker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "AaO297kOiTxD"
   },
   "outputs": [],
   "source": [
    "for m_id in models:\n",
    "  models[m_id].move(secure_worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2nNOeQXjkMOK"
   },
   "source": [
    "## Combine the parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we're collecting each model's parameters and averaging them together. Once the average is obtained, we assign it to our global model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 395
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 322,
     "status": "error",
     "timestamp": 1565377163417,
     "user": {
      "displayName": "Alejandro Aristizabal",
      "photoUrl": "https://lh6.googleusercontent.com/-nfBO6mFL4Mk/AAAAAAAAAAI/AAAAAAAAL1o/9ct9fKu4tjk/s64/photo.jpg",
      "userId": "17089498558649987582"
     },
     "user_tz": 300
    },
    "id": "qQXrzfkQkC4p",
    "outputId": "4a7d321b-9b6f-44cd-8788-57ffd5daab46"
   },
   "outputs": [],
   "source": [
    "# Iterate over every layer of the global model\n",
    "for i, layer in enumerate(list(net.children())):\n",
    "  \n",
    "  # Check if current layer has weights and biases\n",
    "  if len(list(layer.parameters())) != 2:\n",
    "    continue\n",
    "  \n",
    "  # Set variables to store the aggregation\n",
    "  weight = None\n",
    "  bias = None\n",
    "  \n",
    "  # Iterate over every model\n",
    "  for m_id in models:\n",
    "    model = models[m_id]\n",
    "    m_layer = list(model.children())[i]\n",
    "    \n",
    "    # Aggregate current layer's weight\n",
    "    if weight is None:\n",
    "      weight = m_layer.weight\n",
    "    else:\n",
    "      weight += m_layer.weight\n",
    "    \n",
    "    # Aggregate current layer's bias\n",
    "    if bias is None:\n",
    "      bias = m_layer.bias\n",
    "    else:\n",
    "      bias += m_layer.bias\n",
    "      \n",
    "      \n",
    "  # Assign the parameters to our global model\n",
    "  with th.no_grad():\n",
    "    \n",
    "    # Would be nice to do the mean inside secure_worker,\n",
    "    # but float_prec() is not currently supported on pointer tensors \n",
    "    # https://github.com/OpenMined/PySyft/pull/2443\n",
    "    \n",
    "    layer.weight.set_(weight.get().get().float_prec()/len(models))\n",
    "    layer.bias.set_(bias.get().get().float_prec()/len(models))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qtu1b5pno1S7"
   },
   "source": [
    "## Update all local models with out global model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uweKxsVTo41F"
   },
   "source": [
    "Our global model represents a more generalized model. Now we can replace each subject's models with this one. This way, we're obtaining the benefit of using multiple subject's data while retaining privacy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sybAdK7epLjZ"
   },
   "outputs": [],
   "source": [
    "for worker in workers:\n",
    "  models[worker.id] = net.copy().send(worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "RTuVsLwMpTpa"
   },
   "source": [
    "## Check the results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OfhZGjfapWA7"
   },
   "source": [
    "This model should have a good accuracy on every subject. Let's test this hypothesis by running the model locally on our remote testloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 238
    },
    "colab_type": "code",
    "executionInfo": {
     "elapsed": 4374,
     "status": "ok",
     "timestamp": 1565377144070,
     "user": {
      "displayName": "Alejandro Aristizabal",
      "photoUrl": "https://lh6.googleusercontent.com/-nfBO6mFL4Mk/AAAAAAAAAAI/AAAAAAAAL1o/9ct9fKu4tjk/s64/photo.jpg",
      "userId": "17089498558649987582"
     },
     "user_tz": 300
    },
    "id": "fNGSavLBpgTG",
    "outputId": "6a144a93-211b-496b-f83c-19b3f2dc95c8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subject S5.. Test loss: 0.194.. Test accuracy: 0.967\n",
      "Subject S2.. Test loss: 0.401.. Test accuracy: 0.890\n",
      "Subject S3.. Test loss: 0.147.. Test accuracy: 1.000\n",
      "Subject S4.. Test loss: 0.138.. Test accuracy: 0.933\n",
      "Subject S10.. Test loss: 0.226.. Test accuracy: 0.866\n",
      "Subject S11.. Test loss: 0.293.. Test accuracy: 0.915\n",
      "Subject S8.. Test loss: 0.315.. Test accuracy: 0.843\n",
      "Subject S1.. Test loss: 0.268.. Test accuracy: 0.925\n",
      "Subject S6.. Test loss: 0.185.. Test accuracy: 1.000\n",
      "Subject S7.. Test loss: 0.177.. Test accuracy: 0.935\n",
      "Subject S9.. Test loss: 0.108.. Test accuracy: 1.000\n",
      "Subject S13.. Test loss: 0.394.. Test accuracy: 0.883\n",
      "Subject S12.. Test loss: 0.088.. Test accuracy: 1.000\n"
     ]
    }
   ],
   "source": [
    "for worker in workers:\n",
    "  \n",
    "  model = models[worker.id]\n",
    "  testloader = testloaders[worker.id]\n",
    "  \n",
    "  criterion = nn.NLLLoss()\n",
    "  \n",
    "  test_loss = 0\n",
    "  accuracy = 0\n",
    "  model.eval()\n",
    "    \n",
    "  with torch.no_grad():\n",
    "    for inputs, labels in testloader:\n",
    "      inputs, labels = inputs.to(device), labels.to(device)\n",
    "      logps = model.forward(inputs)\n",
    "      test_loss += criterion(logps, labels)\n",
    "        \n",
    "      ps = torch.exp(logps)\n",
    "      top_p, top_class = ps.topk(1, dim=1)\n",
    "      equals = top_class == labels.view(*top_class.shape)\n",
    "      accuracy += torch.mean(equals.float())\n",
    "      \n",
    "  print(f\"Subject {worker.id}.. \"\n",
    "    f\"Test loss: {test_loss.get()/len(testloader):.3f}.. \"\n",
    "    f\"Test accuracy: {accuracy.get()/len(testloader):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Y9N14V2q1TSK"
   },
   "source": [
    "## [Part 4] - Conclusions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using PySyft, we were able to use disjoint private data to train a generalized model without having to compromise or move the data from it's original location. Even though the overall accuracy of our global model is lower than that seen on each individual's model, our model generates a good accuracy for all our subjects. This probably wouldn't happen with any remote model, as they are prone to overfitting on small and low diversity datasets. In a real life scenario, more subjects and more data would be collected, which would result in a more robust and generalized model, and therefore better overall accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done with PySyft version 0.1.23a1\n"
     ]
    }
   ],
   "source": [
    "print(\"Done with PySyft version \" + sy.__version__)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Federated Learning with CNN",
   "provenance": [],
   "version": "0.3.2"
  },
  "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": 1
}