PyTorch/NoteBooks/FL/Homomorphic_Encryption.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "metadata": false
    }
   },
   "source": [
    "## Running FL with secure aggregation using homomorphic encryption\n",
    "\n",
    "This notebook will walk you through how to setup FL with homomorphic encryption (HE). \n",
    "\n",
    "\n",
    "## Prerequisites\n",
    "Before starting this notebook, please make yourself familiar with other FL notebooks in this repo.\n",
    "\n",
    "- (Optional) Look at the introduction Notebook for [Federated Learning with Clara Train SDK](FederatedLearning.ipynb).\n",
    "- (Optional) Look at [Client Notebook](Client.ipynb).\n",
    "- (Optional) Look at [Admin Notebook](Admin.ipynb).\n",
    "- Run [Provisioning Notebook](Provisioning.ipynb) and started the server.\n",
    "\n",
    "Make sure the project.yml used for provision contains these HE related settings:\n",
    "\n",
    "    # homomorphic encryption\n",
    "    he:\n",
    "      lib: tenseal\n",
    "      config:\n",
    "        poly_modulus_degree: 8192\n",
    "        coeff_mod_bit_sizes: [60, 40, 40]\n",
    "        scale_bits: 40\n",
    "        scheme: CKKS\n",
    "        \n",
    "> **Note:** These settings are recommended and should work for most tasks but could be further optimized depending on your specific model architecture and machine learning task. See this [tutorial on the CKKS scheme](https://github.com/OpenMined/TenSEAL/blob/master/tutorials/Tutorial%202%20-%20Working%20with%20Approximate%20Numbers.ipynb) and [benchmarking](https://github.com/OpenMined/TenSEAL/blob/master/tutorials/Tutorial%203%20-%20Benchmarks.ipynb) for more information on different settings."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset \n",
    "\n",
    "##### Option 1 \n",
    "This notebook uses a sample dataset (i.e. a single image volume of the spleen dataset) provided in the package to train a small neural network for a few epochs. \n",
    "This single file is duplicated 32 times for the training set and 9 times for the validation set to mimic the full spleen dataset. \n",
    "\n",
    "##### Option 2  \n",
    "You could do minor changes as recommended in the exercise section in the [Admin notebook](./Admin.ipynb) to train on the spleen segmentation task. The dataset used is Task09_Spleen.tar from \n",
    "the [Medical Segmentation Decathlon](http://medicaldecathlon.com/). \n",
    "Prior to running this notebook the data should be downloaded following \n",
    "the steps in [Data Download Notebook](../../Data_Download.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Disclaimer  \n",
    "We will be training a small networks so that both clients can fit the model on 1 GPU. \n",
    "Training will run for a couple of epochs, in order to show the concepts, we are not targeting accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lets get started\n",
    "In order to learn how FL works with homomorphic encryption (HE) in Clara Train SDK we will first give some background on what homomorphic encryption is and how the MMAR configurations need to be modified to enable it.\n",
    "<br><img src=\"./screenShots/homomorphic_encryption.png\" alt=\"Drawing\" style=\"height: 450px;\"/><br> "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# New HE components\n",
    "\n",
    "We implemented secure aggregation during FL with homomorphic encryption using the [TenSEAL library](https://github.com/OpenMined/TenSEAL) by OpenMined, a convenient wrapper around [Microsoft SEAL](https://github.com/microsoft/SEAL). Both libraries are available as open-source and provide an implementation of [\"Homomorphic encryption for arithmetic of approximate numbers\"](https://eprint.iacr.org/2016/421.pdf), aka the \"CKKS\" scheme, which was proposed as a solution for [encrypted machine learning](https://en.wikipedia.org/wiki/Homomorphic_encryption#Fourth-generation_FHE) and which we use for these FL experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The configuration files in `adminMMAR_HE` use the following new componets that are needed on top or instead of standard FL components used in Clara Train. See the [nvflare docs](https://docs.nvidia.com/clara/clara-train-sdk/federated-learning/apidocs/nvflare.html) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Client-side \n",
    "See `config_fed_client.json`:\n",
    "\n",
    "### `HEModelEncryptor`\n",
    "A filter to the encrypt Shareable object that being sent to the server.\n",
    "\n",
    "```\n",
    "Args:\n",
    "    tenseal_context_file: tenseal context files containing encryption keys and parameters\n",
    "    encrypt_layers: if not specified (None), all layers are being encrypted;\n",
    "                    if list of variable/layer names, only specified variables are encrypted;\n",
    "                    if string containing regular expression (e.g. \"conv\"), only matched variables are being encrypted.\n",
    "    aggregation_weights: dictionary of client aggregation `{\"client1\": 1.0, \"client2\": 2.0, \"client3\": 3.0}`;\n",
    "                         defaults to a weight of 1.0 if not specified.\n",
    "    weigh_by_local_iter: If true, multiply client weights on first before encryption (default: `True` which is recommended for HE)\n",
    "```\n",
    "\n",
    "HE will increase the message sizes when encrypting the model updates of each client. One can choose to not encrypt all layers but specify which layers to encrypt, see argument `encrypt_layers`.\n",
    "\n",
    "To choose the layer names for a given model, one can use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model.0.conv.unit0.conv.weight\n",
      "model.0.conv.unit0.conv.bias\n",
      "model.0.conv.unit0.adn.A.weight\n",
      "model.0.conv.unit1.conv.weight\n",
      "model.0.conv.unit1.conv.bias\n",
      "model.0.conv.unit1.adn.A.weight\n",
      "model.1.submodule.0.conv.unit0.conv.weight\n",
      "model.1.submodule.0.conv.unit0.conv.bias\n",
      "model.1.submodule.0.conv.unit0.adn.A.weight\n",
      "model.1.submodule.0.conv.unit1.conv.weight\n",
      "model.1.submodule.0.conv.unit1.conv.bias\n",
      "model.1.submodule.0.conv.unit1.adn.A.weight\n",
      "model.1.submodule.1.submodule.0.conv.unit0.conv.weight\n",
      "model.1.submodule.1.submodule.0.conv.unit0.conv.bias\n",
      "model.1.submodule.1.submodule.0.conv.unit0.adn.A.weight\n",
      "model.1.submodule.1.submodule.0.conv.unit1.conv.weight\n",
      "model.1.submodule.1.submodule.0.conv.unit1.conv.bias\n",
      "model.1.submodule.1.submodule.0.conv.unit1.adn.A.weight\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.A.weight\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.A.weight\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.A.weight\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.A.weight\n",
      "model.1.submodule.1.submodule.1.submodule.2.0.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.2.0.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.weight\n",
      "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.bias\n",
      "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.A.weight\n",
      "model.1.submodule.1.submodule.2.0.conv.weight\n",
      "model.1.submodule.1.submodule.2.0.conv.bias\n",
      "model.1.submodule.1.submodule.2.1.conv.unit0.conv.weight\n",
      "model.1.submodule.1.submodule.2.1.conv.unit0.conv.bias\n",
      "model.1.submodule.1.submodule.2.1.conv.unit0.adn.A.weight\n",
      "model.1.submodule.2.0.conv.weight\n",
      "model.1.submodule.2.0.conv.bias\n",
      "model.1.submodule.2.1.conv.unit0.conv.weight\n",
      "model.1.submodule.2.1.conv.unit0.conv.bias\n",
      "model.1.submodule.2.1.conv.unit0.adn.A.weight\n",
      "model.2.0.conv.weight\n",
      "model.2.0.conv.bias\n",
      "model.2.1.conv.unit0.conv.weight\n",
      "model.2.1.conv.unit0.conv.bias\n"
     ]
    }
   ],
   "source": [
    "from monai.networks.nets.unet import UNet\n",
    "\n",
    "# use the same configuration as in adminMMAR_HE\n",
    "net = UNet(\n",
    "    dimensions=3,\n",
    "    in_channels=1,\n",
    "    out_channels=2,\n",
    "    channels=[16, 32, 64, 128, 256],\n",
    "    strides=[2, 2, 2, 2],\n",
    "    num_res_units=2    \n",
    ")\n",
    "\n",
    "# here, we only print convolutional layers that we might want to encrypt\n",
    "for key in net.state_dict().keys():\n",
    "    if 'conv' in key:\n",
    "        print(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on this output, our example `config_fed_client.json` chooses three layers to encrypt:\n",
    "\n",
    "```\n",
    "\"outbound_filters\": [\n",
    "  {\n",
    "    \"path\": \"flare.experimental.homomorphic_encryption.he_model_encryptor.HEModelEncryptor\",\n",
    "    \"args\": {\n",
    "      \"encrypt_layers\": [\n",
    "        \"model.0.conv.unit0.conv.weight\",\n",
    "        \"model.1.submodule.1.submodule.1.submodule.2.0.conv.weight\",\n",
    "        \"model.2.1.conv.unit0.conv.weight\"\n",
    "      ],\n",
    "      \"aggregation_weights\": {\n",
    "        \"client1\":  0.4,\n",
    "        \"client2\":  0.6\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "```\n",
    "\n",
    "Using this setting, the client will see an output like this:\n",
    "\n",
    "```\n",
    "2021-05-14 18:23:53,547 - HEModelEncryptor - INFO - Running HE Encryption algorithm on 63 variables\n",
    "2021-05-14 18:23:53,547 - HEModelEncryptor - INFO - Encrypting vars 1 of 63: model.0.conv.unit0.conv.weight with 432 values\n",
    "2021-05-14 18:23:53,566 - HEModelEncryptor - INFO - Encrypting vars 41 of 63: model.1.submodule.1.submodule.1.submodule.2.0.conv.weight with 663552 values\n",
    "WARNING: The input does not fit in a single ciphertext, and some operations will be disabled.\n",
    "The following operations are disabled in this setup: matmul, matmul_plain, enc_matmul_plain, conv2d_im2col.\n",
    "If you need to use those operations, try increasing the poly_modulus parameter, to fit your input.\n",
    "2021-05-14 18:23:54,086 - HEModelEncryptor - INFO - Encrypting vars 62 of 63: model.2.1.conv.unit0.conv.weight with 108 values\n",
    "2021-05-14 18:23:54,089 - HEModelEncryptor - INFO - Encryption time for 664092 of 4806481 params (encrypted value range [-0.0016116723418235779, 0.0016116723418235779]) 0.541804313659668 seconds.\n",
    "```\n",
    "\n",
    "> **Note:** the warning from TenSEAL is expected and can be ignored as these more advanced operations are not used in the `FedAvg` algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEModelDecryptor`\n",
    "A filter to decrypt Shareable object, i.e. the updated global model received from the server.\n",
    "\n",
    "```\n",
    "Args:\n",
    "    tenseal_context_file: tenseal context files containing decryption keys and parameters\n",
    "```\n",
    "\n",
    "> **Note:** The tenseal_context_file for the client will be generated by the provision tool and is part of the startup kit, see [Provisioning](./Provisioning.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEPTModelReaderWriter`\n",
    "\n",
    "This component is used as argument to `ClientTrainer` to reshape the decrypted parameter vectors to the local Pytorch model for training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEEvalDecryptor`\n",
    "Filter to decrypt encrypted Shareable object (i.e. global model(s)) distributed during cross-site validation. Currently, only the global server models are encrypted. Locally best models are shared unencrypted.\n",
    "\n",
    "> **Note:** cross-site validation is optional and a client does not need to participate if not wanted.\n",
    "\n",
    "```\n",
    "Args:\n",
    "    tenseal_context_file: tenseal context files containing decryption keys and parameters\n",
    "    defaults to `False` for use during FL training\n",
    "```\n",
    "\n",
    "The client will see an output like this\n",
    "```\n",
    "2021-05-14 18:23:59,886 - HEModelDecryptor - INFO - Running HE Decryption algorithm 63 variables\n",
    "2021-05-14 18:23:59,886 - HEModelDecryptor - INFO - Decrypting vars 1 of 63: model.0.conv.unit0.conv.weight with 432 values\n",
    "2021-05-14 18:23:59,887 - HEModelDecryptor - INFO - Decrypting vars 41 of 63: model.1.submodule.1.submodule.1.submodule.2.0.conv.weight with 663552 values\n",
    "2021-05-14 18:23:59,976 - HEModelDecryptor - INFO - Decrypting vars 62 of 63: model.2.1.conv.unit0.conv.weight with 108 values\n",
    "2021-05-14 18:23:59,977 - HEModelDecryptor - INFO - Decryption time for 664092 of 664092 params 0.09076642990112305 seconds.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Server-side \n",
    "See `config_fed_server.json`:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEInTimeAccumulateWeightedAggregator`\n",
    "\n",
    "This aggregator can perform federated averaging (i.e. the [`FedAvg`](http://proceedings.mlr.press/v54/mcmahan17a.html) algorithm) in encrypted space. The server doesn't have a key for decryption and only processes the encrypted values sent by the clients.)\n",
    "\n",
    "```\n",
    "Args:\n",
    "    exclude_vars: variable names that should be excluded from aggregation (use regular expression)\n",
    "    aggregation_weights: dictionary of client aggregation `{\"client1\": 1.0, \"client2\": 2.0, \"client3\": 3.0}`;\n",
    "                         defaults to a weight of 1.0 if not specified. Will be ignored if weigh_by_local_iter: False (default for HE)\n",
    "    weigh_by_local_iter: If true, multiply client weights on first in encryption space\n",
    "                         (default: `False` which is recommended for HE, first multiply happens in `HEModelEncryptor`)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEPTFileModelPersistor`\n",
    "\n",
    "This model persistor is used to save the encrypted models on the server."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `HEModelShareableGenerator`\n",
    "\n",
    "ShareableGenerator converts between Shareable and Learnable objects generated with HE. It is used to update the global model weights using the averaged encrypted updates from the clients. The updated global stays encrypted.\n",
    "    \n",
    "```\n",
    "Args:\n",
    "    tenseal_context_file: tenseal context files containing decryption keys and parameters\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output on the server shows how the encrypted layer updates are being aggregated and the global models are updated and saved on the server and the best global model is being saved using `IntimeModelSelectionHandler`.\n",
    "\n",
    "```\n",
    "2021-05-14 18:25:08,354 - FederatedServer - INFO - received Project1_1e41cdaa-e454-432f-a709-fde535d7e03b_1 (55173069 Bytes, 71 seconds)\n",
    "2021-05-14 18:25:08,354 - IntimeModelSelectionHandler - INFO - validation metric 0.05224878713488579 from client client1\n",
    "2021-05-14 18:25:08,354 - HEInTimeAccumulateWeightedAggregator - INFO - Adding contribution from client1.\n",
    "2021-05-14 18:25:08,453 - HEInTimeAccumulateWeightedAggregator - INFO - 3 of 63 layers are encrypted.\n",
    "2021-05-14 18:25:08,454 - HEInTimeAccumulateWeightedAggregator - INFO - Round 1 adding client1 time is 0.09885239601135254 seconds\n",
    "2021-05-14 18:25:13,519 - ServerCrossSiteValManager - INFO - Received best model from client1\n",
    "2021-05-14 18:25:13,543 - ServerCrossSiteValManager - INFO - Client client1 requested models for cross site validation.\n",
    "2021-05-14 18:25:13,543 - ServerCrossSiteValManager - INFO - Sent 0 out of 3 models to client1. Will be asked to wait and retry.\n",
    "2021-05-14 18:25:16,626 - FederatedServer - INFO - received Project1_49ed5224-0962-4683-a20e-11fe16082a88_1 (55172936 Bytes, 79 seconds)\n",
    "2021-05-14 18:25:16,626 - IntimeModelSelectionHandler - INFO - validation metric 0.05224878713488579 from client client2\n",
    "2021-05-14 18:25:16,626 - HEInTimeAccumulateWeightedAggregator - INFO - Adding contribution from client2.\n",
    "2021-05-14 18:25:16,738 - HEInTimeAccumulateWeightedAggregator - INFO - 3 of 63 layers are encrypted.\n",
    "2021-05-14 18:25:16,739 - HEInTimeAccumulateWeightedAggregator - INFO - Round 1 adding client2 time is 0.11211466789245605 seconds\n",
    "2021-05-14 18:25:16,739 - IntimeModelSelectionHandler - INFO - weighted validation metric 0.05224878713488579\n",
    "2021-05-14 18:25:16,739 - IntimeModelSelectionHandler - INFO - new best validation metric at round 1: 0.05224878713488579\n",
    "2021-05-14 18:25:16,739 - FederatedServer - INFO - > aggregating: 1\n",
    "2021-05-14 18:25:16,740 - HEPTFileModelPersistor - INFO - Saving encrypted model on server...\n",
    "2021-05-14 18:25:16,740 - HEPTFileModelPersistor - INFO - 3 of 63 layers are encrypted.\n",
    "2021-05-14 18:25:16,774 - HEPTFileModelPersistor - INFO - Saved encrypted model at /claraDevDay/FL/project1/server/startup/../run_1/mmar_server/models/best_FL_global_model.pt\n",
    "2021-05-14 18:25:16,859 - HEInTimeAccumulateWeightedAggregator - INFO - Aggregated 2 contributions for round 1 time is 0.08475756645202637 seconds\n",
    "2021-05-14 18:25:16,859 - HEInTimeAccumulateWeightedAggregator - INFO - 3 of 63 layers are encrypted.\n",
    "2021-05-14 18:25:16,859 - HEModelShareableGenerator - INFO - HEModelShareableGenerator shareable_to_learnable...\n",
    "2021-05-14 18:25:16,859 - HEModelShareableGenerator - INFO - serialize encrypted model.0.conv.unit0.conv.weight\n",
    "2021-05-14 18:25:16,891 - HEModelShareableGenerator - INFO - serialize encrypted model.1.submodule.1.submodule.1.submodule.2.0.conv.weight\n",
    "2021-05-14 18:25:16,954 - HEModelShareableGenerator - INFO - serialize encrypted model.2.1.conv.unit0.conv.weight\n",
    "2021-05-14 18:25:16,954 - HEModelShareableGenerator - INFO - Updated global model 63 vars with 4806481 params in 0.0953524112701416 seconds\n",
    "2021-05-14 18:25:16,955 - HEPTFileModelPersistor - INFO - Saving encrypted model on server...\n",
    "2021-05-14 18:25:16,955 - HEPTFileModelPersistor - INFO - 3 of 63 layers are encrypted.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running FL experiment with HE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 - Start server, and clients (if they are not already running)\n",
    "Open four terminals in JupyterLab.\n",
    "\n",
    "In the server terminal run:\n",
    "```\n",
    "cd /claraDevDay/FL/project1/server/startup\n",
    "./start.sh\n",
    "```  \n",
    "In the client1 terminal run:\n",
    "```\n",
    "cd /claraDevDay/FL/project1/client1/startup\n",
    "./start.sh\n",
    "```  \n",
    "In the client2 terminal run:\n",
    "```\n",
    "cd /claraDevDay/FL/project1/client2/startup\n",
    "./start.sh\n",
    "```  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 - Starting Admin Shell\n",
    "In the admin terminal, if you haven't already started the admin console you should to admin folder in side your project and run\n",
    "```\n",
    "cd /claraDevDay/FL/project1/admin/startup\n",
    "./fl_admin.sh\n",
    "``` \n",
    "you should see\n",
    "```\n",
    "Admin Server: localhost on port 5000\n",
    "User Name: `\n",
    "```\n",
    "type `admin@admin.com` \n",
    "\n",
    "Admin Server: localhost on port 8003\n",
    "User Name: admin@admin.com\n",
    "\n",
    "Type ? to list commands; type \"? cmdName\" to show usage of a command."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 - Check server/client status\n",
    "type \n",
    "```\n",
    "> check_status server\n",
    "```\n",
    "to see \n",
    "```\n",
    "FL run number has not been set.\n",
    "FL server status: training not started\n",
    "Registered clients: 2 \n",
    "-------------------------------------------------------------------------------------------------\n",
    "| CLIENT NAME | TOKEN                                | LAST ACCEPTED ROUND | CONTRIBUTION COUNT |\n",
    "-------------------------------------------------------------------------------------------------\n",
    "| client1     | f735c245-ce35-4a08-89e0-0292bb053a9c |                     | 0                  |\n",
    "| client2     | e36db52e-2624-4989-855a-28fa195f58e9 |                     | 0                  |\n",
    "-------------------------------------------------------------------------------------------------\n",
    "```\n",
    "To check on clients type \n",
    "```\n",
    "> check_status client\n",
    "```\n",
    "to see \n",
    "```\n",
    "instance:client1 : client name: client1 token: 3c3d2276-c3bf-40c1-bc02-9be84d7c339f     status: training not started\n",
    "instance:client2 : client name: client2 token: 92806548-5515-4977-894e-612900ff8b1b     status: training not started\n",
    "```\n",
    "To check on folder structure \n",
    "\n",
    "```\n",
    "> info\n",
    "```\n",
    "To see\n",
    "```\n",
    "Local Upload Source: /claraDevDay/FL/project1/admin/startup/../transfer\n",
    "Local Download Destination: /claraDevDay/FL/project1/admin/startup/../transfer\n",
    "Server Upload Destination: /claraDevDay/FL/project1/server/startup/../transfer\n",
    "Server Download Source: /claraDevDay/FL/project1/server/startup/../transfer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4- Upload and deploy the MMAR configurations for HE and set FL run number\n",
    "First set a run number (Choose a different one if you don't want to overwrite previous results)\n",
    "```\n",
    "> set_run_number 1\n",
    "```\n",
    "\n",
    "Then, upload the HE MMAR and deploy to server and clients\n",
    "```\n",
    "> upload_folder ../../../adminMMAR_HE\n",
    "> deploy adminMMAR_HE server\n",
    "> deploy adminMMAR_HE client\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 - Start Training\n",
    "Now you can start training by:\n",
    "\n",
    "1. `> start server`\n",
    "2. `> start client`\n",
    "\n",
    "You can check on the status of the training using:\n",
    "\n",
    "3. `> check_status client` or `> check_status server`  to see \n",
    "\n",
    "```\n",
    "FL run number:1\n",
    "FL server status: training started\n",
    "run number:1    start round:0   max round:2     current round:0\n",
    "min_num_clients:2       max_num_clients:100\n",
    "Registered clients: 2 \n",
    "Total number of clients submitted models for current round: 0\n",
    "-------------------------------------------------------------------------------------------------\n",
    "| CLIENT NAME | TOKEN                                | LAST ACCEPTED ROUND | CONTRIBUTION COUNT |\n",
    "-------------------------------------------------------------------------------------------------\n",
    "| client1     | f735c245-ce35-4a08-89e0-0292bb053a9c |                     | 0                  |\n",
    "| client2     | e36db52e-2624-4989-855a-28fa195f58e9 |                     | 0                  |\n",
    "-------------------------------------------------------------------------------------------------\n",
    "```\n",
    "\n",
    "4. get logs from server or clients using `cat server log.txt` or `cat client1 log.txt`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 - Stop Training (if needed ) \n",
    "You could send signals to stop the training if you need to using:\n",
    "- `abort client`\n",
    "- `abort server`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7 - Cross-site validate\n",
    "Once training is completed, you would like to get the validation matrices. See more information about cross-site validation the [Admin notebook](./Admin.ipynb).\n",
    "\n",
    "Run `validate all` to show the cross-site validation results. You could also run `validate source_site target_site` to see the performance of a certain model on a certain site.\n",
    "\n",
    "You should see something like \n",
    "```\n",
    "> validate all\n",
    "{'client2': {'FL_global_model': {'validation': {'val_mean_dice': 0.11059250682592392, 'val_acc': 0.6894438906188668}}, 'client1': {'validation': {'val_mean_dice': 0.12774145603179932, 'val_acc': 0.7208915544895221}}, 'best_FL_global_model': {'validation': {'val_mean_dice': 0.08191516250371933, 'val_acc': 0.625326333805216}}, 'client2': {'validation': {'val_mean_dice': 0.12774255871772766, 'val_acc': 0.7208930085240025}}}, 'client1': {'FL_global_model': {'validation': {'val_mean_dice': 0.1105940118432045, 'val_acc': 0.6894455385246112}}, 'client1': {'validation': {'val_mean_dice': 0.12774255871772766, 'val_acc': 0.7208934608902853}}, 'best_FL_global_model': {'validation': {'val_mean_dice': 0.0819176733493805, 'val_acc': 0.62532652767648}}, 'client2': {'validation': {'val_mean_dice': 0.12774227559566498, 'val_acc': 0.720892297662701}}}}\n",
    "Done [23617 usecs] 2021-05-13 22:04:39.087597\n",
    "``` \n",
    "parsing this json and putting it in a table would look like  \n",
    "\n",
    "Client (val_mean_dice) |  FL_global_model | best_FL_global_model |  client1 |  client2\n",
    ":--- | :--- | :---: | :---: | --- \n",
    "client1  |       0.110594   |           0.081918 | 0.127743 | 0.127742\n",
    "client2  |       0.110593   |           0.081915 | 0.127741 | 0.127743\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8 - Server security\n",
    "\n",
    "To illustrate that the server cannot decrypt the models saved on the server but the client can, we shall execute this small test script.\n",
    "Encrypted models can be loaded using `pickle`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "def load_enc_model(global_model_file):\n",
    "    with open(global_model_file, \"rb\") as f:\n",
    "        model = pickle.load(f)\n",
    "\n",
    "    print(\"model:\", list(model.keys()))\n",
    "    \n",
    "    encrypted_layers = model[\"he_encrypted_layers\"] # holds a True/False boolean indicating whether the layer was encrypted\n",
    "    model = model[\"model\"]  # model state_dict holding the (partially) encrypted model weights\n",
    "\n",
    "    count_encrypted_layers(encrypted_layers)    \n",
    "    \n",
    "    return model, encrypted_layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "> **Note:** The server decryption is expected to fail because the current context of the tensor doesn't hold a secret_key (see output below).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded TenSEAL context from /claraDevDay/FL/project1/server/startup/server_context.tenseal\n",
      "Loaded TenSEAL context from /claraDevDay/FL/project1/client1/startup/client_context.tenseal\n",
      "model: ['model', 'train_conf', 'he_encrypted_layers']\n",
      "3 of 63 layers are encrypted.\n",
      "model.0.conv.unit0.conv.weight is encrypted. Trying to decrypt...\n",
      "<class 'bytes'>\n",
      "Server decryption failed with: the current context of the tensor doesn't hold a secret_key, please provide one as argument!\n",
      "Client decrypted parameters for model.0.conv.unit0.conv.weight:\n",
      "[ 0.07863831  0.15738822  0.12805362 ... -0.07288725  0.05486628\n",
      " -0.0626071 ]\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import tenseal as ts\n",
    "from nvflare.experimental.homomorphic_encryption.homomorphic_encrypt import count_encrypted_layers, load_tenseal_context\n",
    "\n",
    "global_model_file = \"/claraDevDay/FL/project1/server/run_1/mmar_server/models/best_FL_global_model.pt\"\n",
    "server_context_file = \"/claraDevDay/FL/project1/server/startup/server_context.tenseal\"\n",
    "client_context_file = \"/claraDevDay/FL/project1/client1/startup/client_context.tenseal\"\n",
    "\n",
    "# load the server and client TenSEAL context files\n",
    "server_ts_ctx = load_tenseal_context(server_context_file)\n",
    "client_ts_ctx = load_tenseal_context(client_context_file)\n",
    "\n",
    "# load the global model saved on the server\n",
    "model, encrypted_layers = load_enc_model(global_model_file)\n",
    "\n",
    "# try decrypting the first encrypted layer\n",
    "for layer_name in encrypted_layers:\n",
    "    if encrypted_layers[layer_name]:\n",
    "        print(f\"{layer_name} is encrypted. Trying to decrypt...\")\n",
    "        print(type(model[layer_name]))\n",
    "        \n",
    "        try:\n",
    "            # server can deserialize the bytes\n",
    "            ckks_vector = ts.ckks_vector_from(server_ts_ctx, model[layer_name])\n",
    "\n",
    "            # this is supposed to fail with the available server context as it doesn't hold a secret key!\n",
    "            ckks_vector.decrypt()\n",
    "        except Exception as e:\n",
    "            print(f\"Server decryption failed with: {e}!\")\n",
    "            pass\n",
    "        \n",
    "        # However, the client can decrypt using its own TenSEAL context\n",
    "        ckks_vector = ts.ckks_vector_from(client_ts_ctx, model[layer_name])\n",
    "        decrypted_params = ckks_vector.decrypt()\n",
    "        \n",
    "        print(f\"Client decrypted parameters for {layer_name}:\")\n",
    "        np.set_printoptions(threshold=10) # don't show all values\n",
    "        print(np.asarray(decrypted_params))\n",
    "        \n",
    "        break\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10 - Decrypting and validating the global models\n",
    "\n",
    "A client can decrypt the final global models using their TenSEAL context which olds the secret keys.\n",
    "\n",
    "#### First, decrypt the global model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded TenSEAL context from /claraDevDay/FL/project1/client1/startup/client_context.tenseal\n",
      "model: ['model', 'train_conf', 'he_encrypted_layers']\n",
      "3 of 63 layers are encrypted.\n",
      "[-0.19118951273129894 to 0.19245263678326385]\n",
      "[-0.025587799373183125 to 0.02559611276262976]\n",
      "[-0.13731417125904205 to 0.13619750839840944]\n",
      "Saved decrypted model at /claraDevDay/FL/project1/client1/dec_best_FL_global_model.pt\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from monai.networks.nets.unet import UNet\n",
    "\n",
    "global_model_file = \"/claraDevDay/FL/project1/server/run_1/mmar_server/models/best_FL_global_model.pt\"\n",
    "client_context_file = \"/claraDevDay/FL/project1/client1/startup/client_context.tenseal\"\n",
    "dec_global_model_file = \"/claraDevDay/FL/project1/client1/dec_best_FL_global_model.pt\"\n",
    "\n",
    "# load the client TenSEAL context files\n",
    "client_ts_ctx = load_tenseal_context(client_context_file)\n",
    "\n",
    "# load the global model saved on the server\n",
    "model, encrypted_layers = load_enc_model(global_model_file)\n",
    "\n",
    "# configure the model to get the model parameter shapes. Note, this is needed as the encrypted params are just vectors and don't hold the original shapes.\n",
    "net = UNet(\n",
    "  dimensions = 3,\n",
    "  in_channels = 1,\n",
    "  out_channels = 2,\n",
    "  channels = [16, 32, 64, 128, 256],\n",
    "  strides = [2, 2, 2, 2],\n",
    "  num_res_units = 2    \n",
    ")\n",
    "dec_model = net.state_dict()\n",
    "\n",
    "# decrypt encrypted layers and assign to the model\n",
    "for layer_name in dec_model.keys():\n",
    "    if encrypted_layers[layer_name]:\n",
    "        # decrypt using client TenSEAL context\n",
    "        ckks_vector = ts.ckks_vector_from(client_ts_ctx, model[layer_name])\n",
    "        dec_params = ckks_vector.decrypt()\n",
    "        print(f\"[{np.min(dec_params)} to {np.max(dec_params)}]\")\n",
    "    else:\n",
    "        # layer is not encrypted\n",
    "        dec_params = model[layer_name]\n",
    "    dec_model[layer_name] = torch.Tensor(np.reshape(dec_params, dec_model[layer_name].size()))\n",
    "        \n",
    "\n",
    "torch.save(dec_model, dec_global_model_file)\n",
    "print(f\"Saved decrypted model at {dec_global_model_file}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Then, validate the decrypted global model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-05-18 18:52:04,150 - torch.distributed.nn.jit.instantiator - INFO - Created a temporary directory at /tmp/tmps8dahjst\n",
      "2021-05-18 18:52:04,150 - torch.distributed.nn.jit.instantiator - INFO - Writing /tmp/tmps8dahjst/_remote_module_non_sriptable.py\n",
      "========== Validate Config Result ===========\n",
      "Use GPU:  True\n",
      "Multi GPU:  False\n",
      "Automatic Mixed Precision:  Enabled\n",
      "Determinism Evaluation:  Disabled\n",
      "cuDNN BenchMark:  False\n",
      "CUDA Matmul Allow TF32:  True\n",
      "cuDNN Allow TF32:  True\n",
      "Model:  <class 'monai.networks.nets.unet.UNet'>\n",
      "Dataset:  <class 'monai.data.dataset.Dataset'>\n",
      "DataLoader:  <class 'monai.data.dataloader.DataLoader'>\n",
      "Validate Transform #1: <class 'monai.transforms.io.dictionary.LoadImaged'>\n",
      "Validate Transform #2: <class 'monai.transforms.spatial.dictionary.Spacingd'>\n",
      "Validate Transform #3: <class 'monai.transforms.utility.dictionary.EnsureChannelFirstd'>\n",
      "Validate Transform #4: <class 'monai.transforms.intensity.dictionary.ScaleIntensityRanged'>\n",
      "Validate Transform #5: <class 'monai.transforms.croppad.dictionary.CropForegroundd'>\n",
      "Validate Transform #6: <class 'monai.transforms.utility.dictionary.ToTensord'>\n",
      "Validate Handler #1: <class 'monai.handlers.checkpoint_loader.CheckpointLoader'>\n",
      "Validate Handler #2: <class 'monai.handlers.stats_handler.StatsHandler'>\n",
      "Validate Post Transforms #1: <class 'monai.transforms.post.dictionary.Activationsd'>\n",
      "Validate Post Transforms #2: <class 'monai.transforms.post.dictionary.AsDiscreted'>\n",
      "Validate Inferer:  <class 'monai.inferers.inferer.SlidingWindowInferer'>\n",
      "Validate Key Metric:  <class 'monai.handlers.mean_dice.MeanDice'>\n",
      "Validate Additional Metric #val_acc: <class 'ignite.metrics.accuracy.Accuracy'>\n",
      "========== End of Validate Config Result ===========\n",
      "2021-05-18 18:52:06,902 - ignite.engine.engine.SupervisedEvaluator - INFO - Engine run resuming from iteration 0, epoch 0 until 1 epochs\n",
      "2021-05-18 18:52:06,920 - ignite.engine.engine.SupervisedEvaluator - INFO - Restored all variables from /claraDevDay/FL/project1/client1/dec_best_FL_global_model.pt\n",
      "2021-05-18 18:52:08,985 - ignite.engine.engine.SupervisedEvaluator - INFO - Got new best metric of val_mean_dice: 0.05752893164753914\n",
      "2021-05-18 18:52:08,985 - ignite.engine.engine.SupervisedEvaluator - INFO - Epoch[1] Metrics -- val_acc: 0.5468 val_mean_dice: 0.0575 \n",
      "2021-05-18 18:52:08,985 - ignite.engine.engine.SupervisedEvaluator - INFO - Key metric: val_mean_dice best value: 0.05752893164753914 at epoch: 1\n",
      "2021-05-18 18:52:08,985 - ignite.engine.engine.SupervisedEvaluator - INFO - Epoch[1] Complete. Time taken: 00:00:02\n",
      "2021-05-18 18:52:08,986 - ignite.engine.engine.SupervisedEvaluator - INFO - Engine run complete. Time taken: 00:00:02\n",
      "2021-05-18 18:52:09,032 - medl.apps.mmar_conf - INFO - Total Evaluation Time 2.1303601264953613\n"
     ]
    }
   ],
   "source": [
    "!python3 -u  -m medl.apps.evaluate \\\n",
    "    -m \"/claraDevDay/FL/adminMMAR_HE\" \\\n",
    "    -c \"config/config_validation.json\" \\\n",
    "    -e \"config/environment.json\" \\\n",
    "    --set \\\n",
    "    print_conf=True \\\n",
    "    use_gpu=True \\\n",
    "    multi_gpu=False \\\n",
    "    DATA_LIST_KEY=\"validation\" \\\n",
    "    dont_load_ts_model=True \\\n",
    "    dont_load_ckpt_model=False \\\n",
    "    MMAR_CKPT=\"/claraDevDay/FL/project1/client1/dec_best_FL_global_model.pt\" \\\n",
    "    MMAR_EVAL_OUTPUT_PATH=\"/claraDevDay/FL/project1/client1/eval\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11 - Done \n",
    "\n",
    "Congratulations! You have trained and evaluated an FL model using secure aggregation with homomorphic encryption."
   ]
  }
 ],
 "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.8.5"
  },
  "stem_cell": {
   "cell_type": "raw",
   "metadata": {
    "pycharm": {
     "metadata": false
    }
   },
   "source": "<!--- SPDX-License-Identifier: Apache-2.0 -->\n"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}