Task_1/FeTS_Challenge.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FeTS Challenge\n",
    "\n",
    "Contributing Authors (alphabetical order):\n",
    "- Brandon Edwards (Intel)\n",
    "- Patrick Foley (Intel)\n",
    "- Alexey Gruzdev (Intel)\n",
    "- Sarthak Pati (University of Pennsylvania)\n",
    "- Micah Sheller (Intel)\n",
    "- Ilya Trushkin (Intel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '0' # edit according to your system's configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from fets_challenge import run_challenge_experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding custom functionality to the experiment\n",
    "Within this notebook there are **four** functional areas that you can adjust to improve upon the challenge reference code:\n",
    "\n",
    "- [Validation functions](#Custom-Validation-Functions)\n",
    "- [Custom aggregation logic](#Custom-Aggregation-Functions)\n",
    "- [Selection of training hyperparameters by round](#Custom-hyperparameters-for-training)\n",
    "- [Collaborator training selection by round](#Custom-Collaborator-Training-Selection)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experiment logger for your functions\n",
    "The following import allows you to use the same logger used by the experiment framework. This lets you include logging in your functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fets_challenge.experiment import logger"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Validation Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A list of validation function tuples (string_identifier, function_object) should be provided to the validation_functions argument of  run_challenge_experiment to specify the validation functions to perform in addition to the 6 core performance metric functions that will always be computed: (enhancing tumor DICE, tumor core DICE, whole tumor DICE, enhancing tumor hausdorff distance, tumor core hausdorff distance, and whole tumor hausdorff distance). In order to avoid name collision, we have prepended 'performance_evaluation_metric' to each of the string identifiers used for the 6 core permformance metric functions. The string identifier that is included in the 0-index of the tuples will be the assigned name for that metric when it is stored in the aggregator's database. More information about how to use this information for custom aggregation can be found [here](#Using-validation-metrics-for-filtering)\n",
    "\n",
    "Any of the standard PyTorch [validation metrics](https://torchmetrics.readthedocs.io/en/latest/references/modules.html#classification-metrics) can be used to evaluate the model. Any user defined validation functions should conform to the following interface:\n",
    "\n",
    "    def validation_fun_interface(targets, predictions):\n",
    "        \"\"\"validation function interface\n",
    "    \n",
    "        Args:\n",
    "            Targets: numpy array of target values\n",
    "            Predictions: numpy array of predicted values by the model\n",
    "        Returns:\n",
    "            val_score : float\n",
    "        \n",
    "        return val_score\n",
    "        \n",
    "        \n",
    "To add custom metrics to validation that don't conform to the ```(targets, predictions)``` interface, [functool's partial function](https://docs.python.org/3/library/functools.html#functools.partial) can be used to fix a certain number of arguments of a function and generate a new function. For example, we could use [F1 score](https://en.wikipedia.org/wiki/F-score) as a custom metric using the partial function in addition to the exisiting sklearn F1-score metric as follows:\n",
    "```\n",
    "    from functools import partial\n",
    "    from sklearn.metrics import f1_score\n",
    "    validation_functions=[('acc', accuracy), ('f1_score', partial(f1_score, average='macro'))]\n",
    "```\n",
    "        \n",
    "\n",
    "Sensitivity and Specificity are defined below as reference implementations, each performing an average over the enhancing tumor(ET), tumor core(TC), and whole tumor(WT) regions. We utilize a function that takes the float multi-channel model output and multi-label mask and returns binary outputs and masks for each of ET, TC,and WT.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fets_challenge.spec_sens_code import brats_labels\n",
    "\n",
    "\n",
    "def channel_sensitivity(output, target):\n",
    "    # computes TP/P for a single channel \n",
    "\n",
    "    true_positives = np.sum(output * target)\n",
    "    total_positives = np.sum(target)\n",
    "\n",
    "    if total_positives == 0:\n",
    "        score = 1.0\n",
    "    else:\n",
    "        score = true_positives / total_positives\n",
    "    \n",
    "    return score\n",
    "\n",
    "\n",
    "def channel_specificity(output, target):\n",
    "    # computes TN/N for a single channel\n",
    "\n",
    "    true_negatives = np.sum((1 - output) * (1 - target))\n",
    "    total_negatives = np.sum(1 - target)\n",
    "\n",
    "    if total_negatives == 0:\n",
    "        score = 1.0\n",
    "    else:\n",
    "        score = true_negatives / total_negatives\n",
    "        \n",
    "    return score\n",
    "   \n",
    "    \n",
    "def sensitivity(output, target):\n",
    "    \"\"\"\"\n",
    "    Calculates the average sensitivity across all of ET, TC, and WT.\n",
    "    Args:\n",
    "        Targets: numpy array of target values\n",
    "        Predictions: numpy array of predicted values by the model\n",
    "    \"\"\"        \n",
    " \n",
    "    # parsing model output and target into each of ET, TC, and WT arrays\n",
    "    brats_val_data = brats_labels(output=output, target=target)\n",
    "    \n",
    "    outputs = brats_val_data['outputs']\n",
    "    targets = brats_val_data['targets']\n",
    "    \n",
    "    output_enhancing = outputs['ET'] \n",
    "    target_enhancing = targets['ET']\n",
    "\n",
    "    output_core = outputs['TC'] \n",
    "    target_core = targets['TC'] \n",
    "\n",
    "    output_whole = outputs['WT'] \n",
    "    target_whole = targets['WT']\n",
    "\n",
    "    sensitivity_for_enhancing = channel_sensitivity(output=output_enhancing, \n",
    "                                                    target=target_enhancing)\n",
    "\n",
    "    sensitivity_for_core = channel_sensitivity(output=output_core, \n",
    "                                               target=target_core)\n",
    "\n",
    "    sensitivity_for_whole = channel_sensitivity(output=output_whole, \n",
    "                                                target=target_whole)\n",
    "\n",
    "    return (sensitivity_for_enhancing + sensitivity_for_core + sensitivity_for_whole) / 3.0\n",
    "    \n",
    "    \n",
    "def specificity(output, target):\n",
    "    \"\"\"\"\n",
    "    Calculates the average sensitivity across all of ET, TC, and WT.\n",
    "    Args:\n",
    "        Targets: numpy array of target values\n",
    "        Predictions: numpy array of predicted values by the model\n",
    "    \"\"\"  \n",
    "        \n",
    "    # parsing model output and target into each of ET, TC, and WT arrays\n",
    "    brats_val_data = brats_labels(output=output, target=target)\n",
    "    \n",
    "    outputs = brats_val_data['outputs']\n",
    "    targets = brats_val_data['targets']\n",
    "\n",
    "    \n",
    "    output_enhancing = outputs['ET'] \n",
    "    target_enhancing = targets['ET']\n",
    "\n",
    "    output_core = outputs['TC'] \n",
    "    target_core = targets['TC'] \n",
    "\n",
    "    output_whole = outputs['WT'] \n",
    "    target_whole = targets['WT']\n",
    "\n",
    "    specificity_for_enhancing = channel_specificity(output=output_enhancing, \n",
    "                                                    target=target_enhancing)\n",
    "\n",
    "    specificity_for_core = channel_specificity(output=output_core, \n",
    "                                               target=target_core)\n",
    "\n",
    "    specificity_for_whole = channel_specificity(output=output_whole, \n",
    "                                                target=target_whole)\n",
    "\n",
    "    return (specificity_for_enhancing + specificity_for_core + specificity_for_whole) / 3\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting access to historical weights, metrics, and more\n",
    "The **db_iterator** parameter gives full access to all of the tensors and metrics stored by the aggregator. Participants can access these records to create advanced aggregation methods, hyperparameters for training, and novel selection logic for which collaborators should participant in a given training round. See below for details about how data is stored internally and a comprehensive set of examples. \n",
    "\n",
    "## Basic Form\n",
    "Each record yielded by the `db_iterator` contains the following fields:\n",
    "\n",
    "|                      TensorKey                     |   Tensor  |\n",
    "|:--------------------------------------------------:|:---------:|\n",
    "| 'tensor_name', 'origin', 'round', 'report', 'tags' | 'nparray' |\n",
    "\n",
    "All records are internally stored as a numpy array: model weights, metrics, as well as hyperparameters. \n",
    "\n",
    "Detailed field explanation:\n",
    "- **'tensor_name'** (str): The `'tensor_name'` field corresponds to the model layer name (i.e. `'conv2d'`), or the name of the metric that has been reported by a collaborator (i.e. `'accuracy'`). The built-in validation functions used for evaluation of the challenge will be given a prefix of `'challenge_metric_\\*'`. The names that you provide in conjunction with a custom validation metrics to the ```run_challenge_experiment``` function will remain unchanged.  \n",
    "- **'origin'** (str): The origin denotes where the numpy array came from. Possible values are any of the collaborator names (i.e. `'col1'`), or the aggregator.\n",
    "- **'round'** (int): The round that produced the tensor. If your experiment has `N` rounds, possible values are `0->N-1`\n",
    "- **'report'** (boolean): This field is one of the ways that a metric can be denoted; For the purpose of aggregation, this field can be ignored.\n",
    "- **'tags'** (tuple(str)): The tags include unstructured information that can be used to create complex data flows. For example, model layer weights will have the same `'tensor_name'` and `'round'` before and after training, so a tag of `'trained'` is used to denote that the numpy array corresponds to the layer of a locally trained model. This field is also used to capture metric information. For example, `aggregated_model_validation` assigns tags of `'metric'` and `'validate_agg'` to reflect that the metric reported corresponds to the validation score of the latest aggregated model, whereas the tags of `'metric'` and `'validate_local'` are used for metrics produced through validation after training on a collaborator's local data.   \n",
    "- **'nparray'** (numpy array) : This contains the value of the tensor. May contain the model weights, metrics, or hyperparameters as a numpy array.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Note about OpenFL \"tensors\"\n",
    "In order to be ML framework agnostic, OpenFL represents tensors as numpy arrays. Throughout this code, tensor data is represented as numpy arrays (as opposed to torch tensors, for example)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Collaborator Training Selection\n",
    "By default, all collaborators will be selected for training each round, but you can easily add your own logic to select a different set of collaborators based on custom criteria. An example is provided below for selecting a single collaborator on odd rounds that had the fastest training time (`one_collaborator_on_odd_rounds`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a very simple function. Everyone trains every round.\n",
    "def all_collaborators_train(collaborators,\n",
    "                            db_iterator,\n",
    "                            fl_round,\n",
    "                            collaborators_chosen_each_round,\n",
    "                            collaborator_times_per_round):\n",
    "    \"\"\"Chooses which collaborators will train for a given round.\n",
    "    \n",
    "    Args:\n",
    "        collaborators: list of strings of collaborator names\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.  \n",
    "    \"\"\"\n",
    "    return collaborators\n",
    "\n",
    "# this is not a good algorithm, but we include it to demonstrate the following:\n",
    "    # simple use of the logger and of fl_round\n",
    "    # you can search through the \"collaborator_times_per_round\" dictionary to see how long collaborators have been taking\n",
    "    # you can have a subset of collaborators train in a given round\n",
    "def one_collaborator_on_odd_rounds(collaborators,\n",
    "                                   db_iterator,\n",
    "                                   fl_round,\n",
    "                                   collaborators_chosen_each_round,\n",
    "                                   collaborator_times_per_round):\n",
    "    \"\"\"Chooses which collaborators will train for a given round.\n",
    "    \n",
    "    Args:\n",
    "        collaborators: list of strings of collaborator names\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.  \n",
    "    \"\"\"\n",
    "    logger.info(\"one_collaborator_on_odd_rounds called!\")\n",
    "    # on odd rounds, choose the fastest from the previous round\n",
    "    if fl_round % 2 == 1:\n",
    "        training_collaborators = None\n",
    "        fastest_time = np.inf\n",
    "        \n",
    "        # the previous round information will be index [fl_round - 1]\n",
    "        # this information is itself a dictionary of {col: time}\n",
    "        for col, t in collaborator_times_per_round[fl_round - 1].items():\n",
    "            if t < fastest_time:\n",
    "                fastest_time = t\n",
    "                training_collaborators = [col]\n",
    "    else:\n",
    "        training_collaborators = collaborators\n",
    "    return training_collaborators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom hyperparameters for training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can customize the hyper-parameters for the training collaborators at each round. At the start of the round, the experiment loop will invoke your function and set the hyper-parameters for that round based on what you return.\n",
    "\n",
    "The training hyper-parameters for a round are:\n",
    "- **`learning_rate`**: the learning rate value set for the Adam optimizer\n",
    "- **`batches_per_round`**: a flat number of batches each training collaborator will train. Must be an integer or None\n",
    "- **`epochs_per_round`**: the number of epochs each training collaborator will train. Must be a float or None. Partial epochs are allowed, such as 0.5 epochs.\n",
    "\n",
    "Note that exactly one of **`epochs_per_round`** and **`batches_per_round`** must be `\"None\"`. You will get an error message and the experiment will terminate if this is not the case to remind you of this requirement.\n",
    "\n",
    "Your function will receive the typical aggregator state/history information that it can use to make its determination. The function must return a tuple of (`learning_rate`, `epochs_per_round`, `batches_per_round`). For example, if you return:\n",
    "\n",
    "`(1e-4, 2.5, None)`\n",
    "\n",
    "then all collaborators selected based on the [collaborator training selection criteria](#Custom-Collaborator-Training-Selection) will train for `2.5` epochs with a learning rate of `1e-4`.\n",
    "\n",
    "Different hyperparameters can be specified for collaborators for different rounds but they remain the same for all the collaborators that are chosen for that particular round. In simpler words, collaborators can not have different hyperparameters for the same round."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This simple example uses constant hyper-parameters through the experiment\n",
    "def constant_hyper_parameters(collaborators,\n",
    "                              db_iterator,\n",
    "                              fl_round,\n",
    "                              collaborators_chosen_each_round,\n",
    "                              collaborator_times_per_round):\n",
    "    \"\"\"Set the training hyper-parameters for the round.\n",
    "    \n",
    "    Args:\n",
    "        collaborators: list of strings of collaborator names\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.  \n",
    "    Returns:\n",
    "        tuple of (learning_rate, epochs_per_round, batches_per_round). One of epochs_per_round and batches_per_round must be None.\n",
    "    \"\"\"\n",
    "    # these are the hyperparameters used in the May 2021 recent training of the actual FeTS Initiative\n",
    "    # they were tuned using a set of data that UPenn had access to, not on the federation itself\n",
    "    # they worked pretty well for us, but we think you can do better :)\n",
    "    epochs_per_round = 1.0\n",
    "    batches_per_round = None\n",
    "    learning_rate = 5e-5\n",
    "    return (learning_rate, epochs_per_round, batches_per_round)\n",
    "\n",
    "\n",
    "# this example trains less at each round\n",
    "def train_less_each_round(collaborators,\n",
    "                          db_iterator,\n",
    "                          fl_round,\n",
    "                          collaborators_chosen_each_round,\n",
    "                          collaborator_times_per_round):\n",
    "    \"\"\"Set the training hyper-parameters for the round.\n",
    "    \n",
    "    Args:\n",
    "        collaborators: list of strings of collaborator names\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.  \n",
    "    Returns:\n",
    "        tuple of (learning_rate, epochs_per_round, batches_per_round). One of epochs_per_round and batches_per_round must be None.\n",
    "    \"\"\"\n",
    "\n",
    "    # we'll have a constant learning_rate\n",
    "    learning_rate = 5e-5\n",
    "    \n",
    "    # our epochs per round will start at 1.0 and decay by 0.9 for the first 10 rounds\n",
    "    epochs_per_round = 1.0\n",
    "    decay = min(fl_round, 10)\n",
    "    decay = 0.9 ** decay\n",
    "    epochs_per_round *= decay    \n",
    "    \n",
    "    return (learning_rate, epochs_per_round, None)\n",
    "\n",
    "\n",
    "# this example has each institution train the same number of batches\n",
    "def fixed_number_of_batches(collaborators,\n",
    "                            db_iterator,\n",
    "                            fl_round,\n",
    "                            collaborators_chosen_each_round,\n",
    "                            collaborator_times_per_round):\n",
    "    \"\"\"Set the training hyper-parameters for the round.\n",
    "    \n",
    "    Args:\n",
    "        collaborators: list of strings of collaborator names\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.  \n",
    "    Returns:\n",
    "        tuple of (learning_rate, epochs_per_round, batches_per_round). One of epochs_per_round and batches_per_round must be None.\n",
    "    \"\"\"\n",
    "\n",
    "    # we'll have a constant learning_rate\n",
    "    learning_rate = 5e-5\n",
    "    \n",
    "    # instead of a number of epochs, collaborators will train for a number of batches\n",
    "    # this means the number of training batches is irrespective of the data sizes at the institutions\n",
    "    # if the institution has less data than this, they will loop on their data until they have trained\n",
    "    # the correct number of batches\n",
    "    batches_per_round = 16\n",
    "    \n",
    "    # Note that the middle element (epochs_per_round) is now None\n",
    "    return (learning_rate, None, batches_per_round)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Aggregation Functions\n",
    "Standard aggregation methods allow for simple layer-wise combination (via weighted_mean, mean, median, etc.); however, more complex aggregation methods can be supported by evaluating collaborator metrics, weights from prior rounds, etc. OpenFL enables custom aggregation functions via the [**AggregationFunctionInterface**](https://github.com/intel/openfl/blob/fets/openfl/component/aggregation_functions/interface.py). For the challenge, we wrap this interface so we can pass additional simulation state, such as simulated time.\n",
    "\n",
    "[**LocalTensors**](https://github.com/intel/openfl/blob/fets/openfl/utilities/types.py#L13) are named tuples of the form `('collaborator_name', 'tensor', 'collaborator_weight')`. Your custom aggregation function will be passed a list of LocalTensors, which will contain an entry for each collaborator who participated in the prior training round. The [**`db_iterator`**](#Getting-access-to-historical-weights,-metrics,-and-more) can be used to construct complex aggregation methods. A few examples are included below.\n",
    "\n",
    "## db_iterator aggregation examples\n",
    "### Using prior layer weights\n",
    "Here is an example of how to extract layer weights from prior round. The tag is `'aggregated'` indicates this : \n",
    "    \n",
    "    for record in db_iterator:\n",
    "            if (\n",
    "                record['round'] == (fl_round - 1)\n",
    "                and record['tensor_name'] == tensor_name\n",
    "                and 'aggregated' in record['tags']\n",
    "                and 'delta' not in record['tags']\n",
    "               ):\n",
    "                previous_tensor_value = record['nparray']\n",
    "                break\n",
    "\n",
    "### Using validation metrics for filtering\n",
    "\n",
    "    threshold = fl_round * 0.3 + 0.5\n",
    "    metric_name = 'acc'\n",
    "    tags = ('metric','validate_agg')\n",
    "    selected_tensors = []\n",
    "    selected_weights = []\n",
    "    for record in db_iterator:\n",
    "        for local_tensor in local_tensors:\n",
    "            tags = set(tags + [local_tensor.col_name])\n",
    "            if (\n",
    "                tags <= set(record['tags']) \n",
    "                and record['round'] == fl_round\n",
    "                and record['tensor_name'] == metric_name\n",
    "                and record['nparray'] >= threshold\n",
    "            ):\n",
    "                selected_tensors.append(local_tensor.tensor)\n",
    "                selected_weights.append(local_tensor.weight)\n",
    "\n",
    "### A Note about true OpenFL deployments\n",
    "The OpenFL custom aggregation interface does not currently provide timing information, so please note that any solutions that make use of simulated time will need to be adapted to be truly OpenFL compatible in a real federation by using actual `time.time()` calls (or similar) instead of the simulated time.\n",
    "\n",
    "Solutions that use neither **`collaborators_chosen_each_round`** or **`collaborator_times_per_round`** will match the existing OpenFL aggregation customization interface, thus could be used in a real federated deployment using OpenFL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the simple example of weighted FedAVG\n",
    "def weighted_average_aggregation(local_tensors,\n",
    "                                 db_iterator,\n",
    "                                 tensor_name,\n",
    "                                 fl_round,\n",
    "                                 collaborators_chosen_each_round,\n",
    "                                 collaborator_times_per_round):\n",
    "    \"\"\"Aggregate tensors. This aggregator clips all tensor values to the 80th percentile of the absolute values to prevent extreme changes.\n",
    "\n",
    "    Args:\n",
    "        local_tensors(list[openfl.utilities.LocalTensor]): List of local tensors to aggregate.\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        tensor_name: name of the tensor\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.\n",
    "    \"\"\"\n",
    "    # basic weighted fedavg\n",
    "\n",
    "    # here are the tensor values themselves\n",
    "    tensor_values = [t.tensor for t in local_tensors]\n",
    "    \n",
    "    # and the weights (i.e. data sizes)\n",
    "    weight_values = [t.weight for t in local_tensors]\n",
    "    \n",
    "    # so we can just use numpy.average\n",
    "    return np.average(tensor_values, weights=weight_values, axis=0)\n",
    "\n",
    "# here we will clip outliers by clipping deltas to the Nth percentile (e.g. 80th percentile)\n",
    "def clipped_aggregation(local_tensors,\n",
    "                        db_iterator,\n",
    "                        tensor_name,\n",
    "                        fl_round,\n",
    "                        collaborators_chosen_each_round,\n",
    "                        collaborator_times_per_round):\n",
    "    \"\"\"Aggregate tensors. This aggregator clips all tensor values to the 80th percentile of the absolute values to prevent extreme changes.\n",
    "\n",
    "    Args:\n",
    "        local_tensors(list[openfl.utilities.LocalTensor]): List of local tensors to aggregate.\n",
    "        db_iterator: iterator over history of all tensors.\n",
    "            Columns: ['tensor_name', 'round', 'tags', 'nparray']\n",
    "        tensor_name: name of the tensor\n",
    "        fl_round: round number\n",
    "        collaborators_chosen_each_round: a dictionary of {round: list of collaborators}. Each list indicates which collaborators trained in that given round.\n",
    "        collaborator_times_per_round: a dictionary of {round: {collaborator: total_time_taken_in_round}}.\n",
    "    \"\"\"\n",
    "    # the percentile we will clip to\n",
    "    clip_to_percentile = 80\n",
    "    \n",
    "    # first, we need to determine how much each local update has changed the tensor from the previous value\n",
    "    # we'll use the db_iterator to find the previous round's value for this tensor\n",
    "    previous_tensor_value = None\n",
    "    for record in db_iterator:\n",
    "        if (\n",
    "            record['round'] == (fl_round - 1)\n",
    "            and record['tensor_name'] == tensor_name\n",
    "            and 'aggregated' in record['tags']\n",
    "            and 'delta' not in record['tags']\n",
    "           ):\n",
    "            previous_tensor_value = record['nparray']\n",
    "            break\n",
    "       \n",
    "    # if we have no previous tensor_value, we won't actually clip\n",
    "    if previous_tensor_value is None:\n",
    "        clipped_tensors = [t.tensor for t in local_tensors]\n",
    "    # otherwise, we will use clipped deltas\n",
    "    else:\n",
    "        # compute the deltas\n",
    "        deltas = [t.tensor - previous_tensor_value for t in local_tensors]\n",
    "    \n",
    "        # concatenate all the deltas\n",
    "        all_deltas = np.concatenate(deltas)\n",
    "        \n",
    "        # take the absolute value\n",
    "        abs_deltas = np.abs(all_deltas)\n",
    "        \n",
    "        # finally, get the 80th percentile\n",
    "        clip_value = np.percentile(abs_deltas, clip_to_percentile)\n",
    "        \n",
    "        # let's log what we're clipping to\n",
    "        logger.info(\"Clipping tensor {} to value {}\".format(tensor_name, clip_value))\n",
    "    \n",
    "        # now we can compute our clipped tensors\n",
    "        clipped_tensors = []\n",
    "        for delta, t in zip(deltas, local_tensors):\n",
    "            new_tensor = previous_tensor_value + np.clip(delta, -1 * clip_value, clip_value)\n",
    "            clipped_tensors.append(new_tensor)\n",
    "        \n",
    "    # get an array of weight values for the weighted average\n",
    "    weights = [t.weight for t in local_tensors]\n",
    "\n",
    "    # return the weighted average of the clipped tensors\n",
    "    return np.average(clipped_tensors, weights=weights, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running the Experiment\n",
    "\n",
    "```run_challenge_experiment``` is singular interface where your custom methods can be passed.\n",
    "\n",
    "- ```aggregation_function```, ```choose_training_collaborators```, ```training_hyper_parameters_for_round```, and ```validation_functions``` correspond to the [this list](#Custom-hyperparameters-for-training) of configurable functions \n",
    "described within this notebook.\n",
    "- ```validation_functions``` should be a list of tuples, enabling you add multiple additional validation functions. The tuples should be `(name, function)`, where `name` is the string that will be associated with the metric in the `tensor_db`, and `function` is the python function you implemented above. It can be an empty list if you do not wish to add additional validation functions.\n",
    "- ```institution_split_csv_filename``` : Describes how the data should be split between all collaborators. Extended documentation about configuring the splits in the ```institution_split_csv_filename``` parameter can be found in the [README.md](https://github.com/FETS-AI/Challenge/blob/main/Task_1/README.md). \n",
    "- ```db_store_rounds``` : This parameter determines how long metrics and weights should be stored by the aggregator before being deleted. Providing a value of `-1` will result in all historical data being retained, but memory usage will likely increase.\n",
    "- ```rounds_to_train``` : Defines how many rounds will occur in the experiment\n",
    "- ```device``` : Which device to use for training and validation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the experiment\n",
    "Now that we've defined our custom functions, the last thing to do is to configure the experiment. The following cell shows the various settings you can change in your experiment.\n",
    "\n",
    "Note that ```rounds_to_train``` can be set as high as you want. However, the experiment will exit once the simulated time value exceeds 1 week of simulated time, or if the specified number of rounds has completed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change any of these you wish to your custom functions. You may leave defaults if you wish.\n",
    "aggregation_function = weighted_average_aggregation\n",
    "choose_training_collaborators = all_collaborators_train\n",
    "training_hyper_parameters_for_round = constant_hyper_parameters\n",
    "validation_functions = [('sensitivity', sensitivity), ('specificity', specificity)]\n",
    "\n",
    "# We encourage participants to experiment with partitioning_1 and partitioning_2, as well as to create\n",
    "# other partitionings to test your changes for generalization to multiple partitionings.\n",
    "institution_split_csv_filename = 'partitioning_1.csv'\n",
    "\n",
    "# change this to point to the parent directory of the data\n",
    "brats_training_data_parent_dir = '/raid/datasets/FeTS21/MICCAI_FeTS2021_TrainingData'\n",
    "\n",
    "# increase this if you need a longer history for your algorithms\n",
    "# decrease this if you need to reduce system RAM consumption\n",
    "db_store_rounds = 5\n",
    "\n",
    "# this is passed to PyTorch, so set it accordingly for your system\n",
    "device = 'cuda'\n",
    "\n",
    "# you'll want to increase this most likely. You can set it as high as you like, \n",
    "# however, the experiment will exit once the simulated time exceeds one week. \n",
    "rounds_to_train = 5\n",
    "\n",
    "# (bool) Determines whether checkpoints should be saved during the experiment. \n",
    "# The checkpoints can grow quite large (5-10GB) so only the latest will be saved when this parameter is enabled\n",
    "save_checkpoints = True\n",
    "\n",
    "# path to previous checkpoint folder for experiment that was stopped before completion. \n",
    "# Checkpoints are stored in ~/.local/workspace/checkpoint, and you should provide the experiment directory \n",
    "# relative to this path (i.e. 'experiment_1'). Please note that if you restore from a checkpoint, \n",
    "# and save checkpoint is set to True, then the checkpoint you restore from will be subsequently overwritten.\n",
    "# restore_from_checkpoint_folder = 'experiment_1'\n",
    "restore_from_checkpoint_folder = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the scores are returned in a Pandas dataframe\n",
    "scores_dataframe = run_challenge_experiment(\n",
    "    aggregation_function=aggregation_function,\n",
    "    choose_training_collaborators=choose_training_collaborators,\n",
    "    training_hyper_parameters_for_round=training_hyper_parameters_for_round,\n",
    "    validation_functions=validation_functions,\n",
    "    institution_split_csv_filename=institution_split_csv_filename,\n",
    "    brats_training_data_parent_dir=brats_training_data_parent_dir,\n",
    "    db_store_rounds=db_store_rounds,\n",
    "    rounds_to_train=rounds_to_train,\n",
    "    device=device,\n",
    "    save_checkpoints=save_checkpoints,\n",
    "    restore_from_checkpoint_folder = restore_from_checkpoint_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scores_dataframe"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "openfl",
   "language": "python",
   "name": "openfl"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}