diff --git a/examples/splitnn/Dual_headed_SplitNN_Config1.ipynb b/examples/splitnn/Dual_headed_SplitNN_Config1.ipynb
new file mode 100644
index 0000000..65ba2e5
--- /dev/null
+++ b/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",
+ "What is Vertically Partitioned Data? 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",
+ "Why use Partitioned Data? 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",
+ "DescriptionThis 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",
+ "![](\"images/config_1.png\")
\n",
+ "\n",
+ "In this tutorial, we demonstrate the SplitNN architecture with 2 segments[[1](https://arxiv.org/abs/1812.00564)].This time:\n",
+ "\n",
+ "- $Client_{1}$\n",
+ " - Has Model Segment 1\n",
+ " - Has the handwritten images segment 1\n",
+ "- $Client_{2}$\n",
+ " - Has model Segment 1\n",
+ " - Has the handwritten images segment 2\n",
+ "- $Server$ \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
+}
diff --git a/examples/splitnn/Dual_headed_SplitNN_Config2.ipynb b/examples/splitnn/Dual_headed_SplitNN_Config2.ipynb
new file mode 100644
index 0000000..5a45e56
--- /dev/null
+++ b/examples/splitnn/Dual_headed_SplitNN_Config2.ipynb
@@ -0,0 +1,359 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "epochs= 25"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# SplitNN for Vertically Partitioned Data\n",
+ "\n",
+ "What is Vertically Partitioned Data? 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",
+ "Why use Partitioned Data? 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",
+ "DescriptionThis 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",
+ "![](\"images/config-2.jpg\")
\n",
+ "\n",
+ "In this tutorial, we demonstrate the SplitNN architecture with 2 segments[[1](https://arxiv.org/abs/1812.00564)].This time:\n",
+ "\n",
+ "- $Client_{1}$\n",
+ " - Has Model Segment 1\n",
+ " - Has the handwritten images segment 1\n",
+ "- $Client_{2}$\n",
+ " - Has model Segment 1\n",
+ " - Has the handwritten images segment 2\n",
+ "- $Server$ \n",
+ " - Has Model Segment 2\n",
+ "- $Label Owner$ \n",
+ " - Has Model Segment 3\n",
+ " - Has the image labels\n",
+ " \n",
+ " \n",
+ "Author:\n",
+ "- Abbas Ismail - github：[@abbas5253](https://github.com/abbas5253)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "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": [
+ "class SplitNN(torch.nn.Module):\n",
+ " \n",
+ " def __init__(self, models, optimizers, data_owner, label_owner, server):\n",
+ " \n",
+ " self.label_owner = label_owner\n",
+ " self.data_owners = data_owner\n",
+ " self.optimizers = optimizers\n",
+ " self.models = models\n",
+ " self.server = server\n",
+ " super().__init__()\n",
+ " \n",
+ " \n",
+ " def forward(self, data_pointer):\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 move outputs to server\n",
+ " for owner in self.data_owners:\n",
+ " client_output[owner.id] = self.models[owner.id](data_pointer[owner.id].reshape([-1, 14*28]))\n",
+ " remote_outputs.append(\n",
+ " client_output[owner.id].move(self.server).requires_grad_()\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",
+ " server_output = self.models[\"server\"](server_input)\n",
+ " \n",
+ " # move to label_owner\n",
+ " server_output = server_output.move(self.label_owner).requires_grad_()\n",
+ " \n",
+ " #pass through the back segment and return prediction\n",
+ " pred = models[self.label_owner.id](server_output)\n",
+ " \n",
+ " return pred \n",
+ " \n",
+ " def zero_grads(self):\n",
+ " for opt in self.optimizers:\n",
+ " opt.zero_grad()\n",
+ " \n",
+ " def step(self):\n",
+ " for opt in self.optimizers:\n",
+ " opt.step()\n",
+ " \n",
+ " def train(self):\n",
+ " for loc in self.models.keys():\n",
+ " for i in range(len(self.models[loc])):\n",
+ " self.models[loc][i].train()\n",
+ " \n",
+ " def eval(self):\n",
+ " for loc in self.models.keys():\n",
+ " for i in range(len(self.models[loc])):\n",
+ " self.models[loc][i].eval()\n",
+ " \n",
+ "# @property\n",
+ "# def location(self):\n",
+ "# return self.models[0].location if self.models and len(self.models) else None"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Data preprocessing\n",
+ "transform = transforms.Compose([transforms.ToTensor(),\n",
+ " transforms.Normalize((0.5,), (0.5,)),\n",
+ " ])\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",
+ "label_owner = sy.VirtualWorker(hook, id=\"label_owner\")\n",
+ "# label_owner = client_1\n",
+ "\n",
+ "data_owners = (client_1, client_2)\n",
+ "model_locations = [client_1, client_2, server, label_owner]\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": 7,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "input_size= [28*14, 28*14]\n",
+ "hidden_sizes= {\"client_1\": [32], \"client_2\":[32], \"server\":[64, 128]}\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",
+ " ),\n",
+ " \n",
+ " \n",
+ " \"client_2\" : nn.Sequential(\n",
+ " nn.Linear(input_size[1], hidden_sizes[\"client_2\"][0]),\n",
+ " nn.ReLU(),\n",
+ " ),\n",
+ " \n",
+ " \"server\": nn.Sequential(\n",
+ " nn.Linear(hidden_sizes[\"server\"][0], hidden_sizes[\"server\"][1]),\n",
+ " nn.ReLU(),\n",
+ " ),\n",
+ " \n",
+ " \"label_owner\": nn.Sequential(\n",
+ " nn.Linear(128, 64),\n",
+ " nn.ReLU(),\n",
+ " nn.Linear(64, 10),\n",
+ " nn.LogSoftmax(dim=1)\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": 8,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def train(x, target, splitNN):\n",
+ " \n",
+ " #1) Zero our grads\n",
+ " splitNN.zero_grads()\n",
+ " \n",
+ " #2) Make a prediction\n",
+ " pred = splitNN.forward(x)\n",
+ " \n",
+ " #3) Figure out how much we missed by\n",
+ " criterion = nn.NLLLoss()\n",
+ " loss = criterion(pred, target.reshape(-1, 64)[0])\n",
+ " \n",
+ " #4) Backprop the loss on the end layer\n",
+ " loss.backward()\n",
+ " \n",
+ " #5) Feed Gradients backward through the network\n",
+ " #splitNN.backward()\n",
+ " \n",
+ " #6) Change the weights\n",
+ " splitNN.step()\n",
+ " \n",
+ " return loss.detach().get()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 9,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Epoch 0 - Training loss: 2.1097309589385986\n",
+ "Epoch 1 - Training loss: 1.4561858177185059\n",
+ "Epoch 2 - Training loss: 1.1583361625671387\n",
+ "Epoch 3 - Training loss: 1.0346523523330688\n",
+ "Epoch 4 - Training loss: 0.9617491364479065\n",
+ "Epoch 5 - Training loss: 0.9116027355194092\n",
+ "Epoch 6 - Training loss: 0.8750373125076294\n",
+ "Epoch 7 - Training loss: 0.8472383618354797\n",
+ "Epoch 8 - Training loss: 0.825141966342926\n",
+ "Epoch 9 - Training loss: 0.806628406047821\n",
+ "Epoch 10 - Training loss: 0.7905024290084839\n",
+ "Epoch 11 - Training loss: 0.7761190533638\n",
+ "Epoch 12 - Training loss: 0.7629939913749695\n",
+ "Epoch 13 - Training loss: 0.7508707046508789\n",
+ "Epoch 14 - Training loss: 0.739571750164032\n",
+ "Epoch 15 - Training loss: 0.7288739681243896\n",
+ "Epoch 16 - Training loss: 0.7187440395355225\n",
+ "Epoch 17 - Training loss: 0.7090423703193665\n",
+ "Epoch 18 - Training loss: 0.6996704339981079\n",
+ "Epoch 19 - Training loss: 0.6906035542488098\n",
+ "Epoch 20 - Training loss: 0.6819005012512207\n",
+ "Epoch 21 - Training loss: 0.6735816597938538\n",
+ "Epoch 22 - Training loss: 0.6655935049057007\n",
+ "Epoch 23 - Training loss: 0.6579392552375793\n",
+ "Epoch 24 - Training loss: 0.65062016248703\n"
+ ]
+ }
+ ],
+ "source": [
+ "splitnn = SplitNN(models, optimizers, data_owners, label_owner, server)\n",
+ "\n",
+ "for i in range(epochs):\n",
+ " \n",
+ " running_loss = 0\n",
+ " splitnn.train()\n",
+ " for images, labels in distributed_trainloader:\n",
+ " labels = labels.send(label_owner)\n",
+ " loss = train(images, labels, splitnn)\n",
+ " running_loss += loss\n",
+ "\n",
+ " else:\n",
+ " print(\"Epoch {} - Training loss: {}\".format(i, running_loss/len(trainloader)))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "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 = splitnn.forward(data_ptr).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": 11,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Train set: Accuracy 47605/59968 (79%)\n",
+ "Test set: Accuracy 8009/9984 (80%)\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
+}
diff --git a/examples/splitnn/images/config-2.jpg b/examples/splitnn/images/config-2.jpg
new file mode 100644
index 0000000..9bf6ef4
Binary files /dev/null and b/examples/splitnn/images/config-2.jpg differ
diff --git a/examples/splitnn/images/config_1.png b/examples/splitnn/images/config_1.png
new file mode 100644
index 0000000..1762f47
Binary files /dev/null and b/examples/splitnn/images/config_1.png differ