Dual-Headed SplitNN Config1 and Config2

examples/splitnn/Dual_headed_SplitNN_Config1.ipynb

@@ -0,0 +1,327 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "epochs= 25"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# SplitNN for Vertically Partitioned Data\n",
+    "\n",
+    "<b>What is Vertically Partitioned Data?</b> Data is said to be vertically partitioned when several organizations own different attributes or modalities of information for the same set of entities.\n",
+    "\n",
+    "<b>Why use Partitioned Data?</b> Partition allows for orgnizations holding different modalities of data to learn distributed models without data sharing. Partitioning scheme is traditionally used to reduce the size of data by splitting and distribute to each client.\n",
+    " \n",
+    "<b>Description</b>This configuration allows for multiple clients holding different modalities of data to learn distributed models without data sharing. As a concrete example we walkthrough the case where radiology centers collaborate with pathology test centers and a server for disease diagnosis. Radiology centers holding imaging data modalities train a partial model upto the cut layer. In the same way the pathology test center having patient test results trains a partial model upto its own cut layer. The outputs at the cut layer from both these centers are then concatenated and sent to the disease diagnosis server that trains the rest of the model. This process is continued back and forth to complete the forward and backward propagations in order to train the distributed deep learning model without sharing each others raw data. In this tutorial, we split a single flatten image into two segments to mimic different modalities of data, you can also split it into arbitrary number.\n",
+    "\n",
+    "<img src=\"images/config_1.png\" width=\"40%\">\n",
+    "\n",
+    "In this tutorial, we demonstrate the SplitNN architecture with 2 segments[[1](https://arxiv.org/abs/1812.00564)].This time:\n",
+    "\n",
+    "- <b>$Client_{1}$</b>\n",
+    "    - Has Model Segment 1\n",
+    "    - Has the handwritten images segment 1\n",
+    "- <b>$Client_{2}$</b>\n",
+    "    - Has model Segment 1\n",
+    "    - Has the handwritten images segment 2\n",
+    "- <b>$Server$</b> \n",
+    "    - Has Model Segment 2\n",
+    "    - Has the image labels\n",
+    "    \n",
+    "Author:\n",
+    "- Abbas Ismail - github：[@abbas5253](https://github.com/abbas5253)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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 '/home/ab_53/miniconda3/envs/PySyft/lib/python3.7/site-packages/tf_encrypted/operations/secure_random/secure_random_module_tf_1.15.3.so'\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "WARNING:tensorflow:From /home/ab_53/miniconda3/envs/PySyft/lib/python3.7/site-packages/tf_encrypted/session.py:24: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "from torchvision import datasets, transforms\n",
+    "from torch import nn, optim\n",
+    "import syft as sy\n",
+    "import numpy as np\n",
+    "\n",
+    "hook = sy.TorchHook(torch)\n",
+    "\n",
+    "from distribute_data import Distribute_MNIST"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Data preprocessing\n",
+    "transform = transforms.Compose([transforms.ToTensor(),\n",
+    "                              transforms.Normalize((0.5,), (0.5,)),\n",
+    "                              ])\n",
+    "trainset = datasets.MNIST('mnist', download=True, train=True, transform=transform)\n",
+    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)\n",
+    "\n",
+    "# create some workers\n",
+    "client_1 = sy.VirtualWorker(hook, id=\"client_1\")\n",
+    "client_2 = sy.VirtualWorker(hook, id=\"client_2\")\n",
+    "\n",
+    "server = sy.VirtualWorker(hook, id= \"server\") \n",
+    "\n",
+    "data_owners = (client_1, client_2)\n",
+    "model_locations = [client_1, client_2, server]\n",
+    "\n",
+    "#Split each image and send one part to client_1, and other to client_2\n",
+    "distributed_trainloader = Distribute_MNIST(data_owners=data_owners, data_loader=trainloader)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_size= [28*14, 28*14]\n",
+    "hidden_sizes= {\"client_1\": [32, 64], \"client_2\":[32, 64], \"server\":[128, 64]}\n",
+    "\n",
+    "#create model segment for each worker\n",
+    "models = {\n",
+    "    \"client_1\": nn.Sequential(\n",
+    "                nn.Linear(input_size[0], hidden_sizes[\"client_1\"][0]),\n",
+    "                nn.ReLU(),\n",
+    "                nn.Linear(hidden_sizes[\"client_1\"][0], hidden_sizes[\"client_1\"][1]),\n",
+    "                nn.ReLU(),\n",
+    "    ),\n",
+    "    \"client_2\":  nn.Sequential(\n",
+    "                nn.Linear(input_size[1], hidden_sizes[\"client_2\"][0]),\n",
+    "                nn.ReLU(),\n",
+    "                nn.Linear(hidden_sizes[\"client_2\"][0], hidden_sizes[\"client_2\"][1]),\n",
+    "                nn.ReLU(),\n",
+    "    ),\n",
+    "    \"server\": nn.Sequential(\n",
+    "                nn.Linear(hidden_sizes[\"server\"][0], hidden_sizes[\"server\"][1]),\n",
+    "                nn.ReLU(),\n",
+    "                nn.Linear(hidden_sizes[\"server\"][1], 10),\n",
+    "                nn.LogSoftmax(dim=1)\n",
+    "    )\n",
+    "}\n",
+    "\n",
+    "\n",
+    "\n",
+    "# Create optimisers for each segment and link to their segment\n",
+    "optimizers = [\n",
+    "    optim.SGD(models[location.id].parameters(), lr=0.05,)\n",
+    "    for location in model_locations\n",
+    "]\n",
+    "\n",
+    "\n",
+    "#send model segement to each client and server\n",
+    "for location in model_locations:\n",
+    "    models[location.id].send(location)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def predict(data_pointer, models, data_owners, server):\n",
+    "    \n",
+    "    #individual client's output upto their respective cut layer\n",
+    "    client_output = {}\n",
+    "    \n",
+    "    #outputs that is moved to server and subjected to concatenate for server input\n",
+    "    remote_outputs = []\n",
+    "    \n",
+    "    #iterate over each client and pass thier inputs to respective model segment and send outputs to server\n",
+    "    for owner in data_owners:\n",
+    "        client_output[owner.id] = models[owner.id](data_pointer[owner.id].reshape([-1, 14*28]))\n",
+    "        remote_outputs.append(\n",
+    "            client_output[owner.id].move(server)\n",
+    "        )\n",
+    "    \n",
+    "    #concat outputs from all clients at server's location\n",
+    "    server_input = torch.cat(remote_outputs, 1)\n",
+    "    \n",
+    "    #pass concatenated output from server's model segment\n",
+    "    pred = models[\"server\"](server_input)\n",
+    "    \n",
+    "    return pred"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train(data_pointer, target, data_owners, models, optimizers, server):\n",
+    "    \n",
+    "    #make grads zero\n",
+    "    for opt in optimizers:\n",
+    "        opt.zero_grad()\n",
+    "    \n",
+    "    #predict the output\n",
+    "    pred = predict(data_pointer, models, data_owners, server)\n",
+    "    \n",
+    "    #calculate loss\n",
+    "    criterion = nn.NLLLoss()\n",
+    "    loss = criterion(pred, target.reshape(-1, 64)[0])\n",
+    "    \n",
+    "    #backpropagate\n",
+    "    loss.backward()\n",
+    "    \n",
+    "    #optimization step\n",
+    "    for opt in optimizers:\n",
+    "        opt.step()\n",
+    "        \n",
+    "    return loss.detach().get()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0 - Training loss: 2.074364185333252\n",
+      "Epoch 1 - Training loss: 1.4119279384613037\n",
+      "Epoch 2 - Training loss: 1.0729103088378906\n",
+      "Epoch 3 - Training loss: 0.9095255136489868\n",
+      "Epoch 4 - Training loss: 0.8234174251556396\n",
+      "Epoch 5 - Training loss: 0.7714390754699707\n",
+      "Epoch 6 - Training loss: 0.7359529733657837\n",
+      "Epoch 7 - Training loss: 0.7096782922744751\n",
+      "Epoch 8 - Training loss: 0.6890704035758972\n",
+      "Epoch 9 - Training loss: 0.6721634268760681\n",
+      "Epoch 10 - Training loss: 0.65777188539505\n",
+      "Epoch 11 - Training loss: 0.6452123522758484\n",
+      "Epoch 12 - Training loss: 0.6340039968490601\n",
+      "Epoch 13 - Training loss: 0.6238173246383667\n",
+      "Epoch 14 - Training loss: 0.6144205331802368\n",
+      "Epoch 15 - Training loss: 0.6057182550430298\n",
+      "Epoch 16 - Training loss: 0.5975443124771118\n",
+      "Epoch 17 - Training loss: 0.5897234082221985\n",
+      "Epoch 18 - Training loss: 0.5822254419326782\n",
+      "Epoch 19 - Training loss: 0.5749425888061523\n",
+      "Epoch 20 - Training loss: 0.5678632259368896\n",
+      "Epoch 21 - Training loss: 0.5609995722770691\n",
+      "Epoch 22 - Training loss: 0.5543981790542603\n",
+      "Epoch 23 - Training loss: 0.5480557084083557\n",
+      "Epoch 24 - Training loss: 0.5418789386749268\n"
+     ]
+    }
+   ],
+   "source": [
+    "#training\n",
+    "for i in range(epochs):\n",
+    "    running_loss = 0\n",
+    "    \n",
+    "    #iterate over each datapoints \n",
+    "    for data_ptr, label in distributed_trainloader:\n",
+    "        \n",
+    "        #send labels to server's location for training\n",
+    "        label = label.send(server)\n",
+    "        \n",
+    "        loss = train(data_ptr, label, data_owners, models, optimizers, server)\n",
+    "        running_loss += loss\n",
+    "\n",
+    "    else:\n",
+    "        print(\"Epoch {} - Training loss: {}\".format(i, running_loss/len(trainloader)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test(model, dataloader, dataset_name):\n",
+    "    correct = 0\n",
+    "    with torch.no_grad():\n",
+    "        for data_ptr, label in dataloader:\n",
+    "            output = predict(data_ptr, models, data_owners, server).get()\n",
+    "            pred = output.max(1, keepdim=True)[1]\n",
+    "            correct += pred.eq(label.data.view_as(pred)).sum()\n",
+    "\n",
+    "    print(\"{}: Accuracy {}/{} ({:.0f}%)\".format(dataset_name, \n",
+    "                                                correct,\n",
+    "                                                len(dataloader)* 64, \n",
+    "                                                100. * correct / (len(dataloader) * 64) ))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Train set: Accuracy 49885/59968 (83%)\n",
+      "Test set: Accuracy 8290/9984 (83%)\n"
+     ]
+    }
+   ],
+   "source": [
+    "#prepare and distribute test dataset\n",
+    "testset = datasets.MNIST('mnist', download=True, train=False, transform=transform)\n",
+    "testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)\n",
+    "distributed_testloader = Distribute_MNIST(data_owners=data_owners, data_loader=testloader)\n",
+    "\n",
+    "#Accuracy on train and test sets\n",
+    "test(models, distributed_trainloader, \"Train set\")\n",
+    "test(models, distributed_testloader, \"Test set\")"
+   ]
+  }
+ ],
+ "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.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}