06-Federated-learning-Experiment-with-PyTorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ecd52778-c827-4909-b4ed-aca9556e956a",
   "metadata": {},
   "source": [
    "### Preamble\n",
    "> *Welcome to the final part - you're officially* ***unstoppable*** 💫 💪\n",
    ">\n",
    "> **You're amazing!** Let's wrap this journey up in style!\n",
    ">\n",
    "> *Valerio*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cad9e42-c449-4487-ba66-0f379a25e4fc",
   "metadata": {},
   "source": [
    "# Federated learning with PySyft and PyTorch\n",
    "\n",
    "In this notebook we will learn how to run a Federated Learning experiment, using **PyTorch** and **PySyft**."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67e95d51-313e-4ddb-9340-5e3de0e82bc0",
   "metadata": {},
   "source": [
    "## Step 1. Login to datasites as **External Researcher**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e167121f-90cc-4bbc-8e15-0e8e56160302",
   "metadata": {},
   "source": [
    "⚠️ First verify that the Datasites are already running. If needed, launch the following command in a new terminal session:\n",
    "\n",
    "```bash\n",
    "$ python launch_datasites.py\n",
    "```\n",
    "\n",
    "**Note**: In Jupyter Lab, you can open a new terminal session via `File >> New >> Terminal`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f3a3ee2-7a98-427c-8098-30d724e43be3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import syft as sy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1283e48c-9d97-42f4-a905-64436ed24a46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasites import DATASITE_URLS\n",
    "\n",
    "datasites = {}\n",
    "for name, url in DATASITE_URLS.items():\n",
    "    datasites[name] = sy.login(url=url, email=\"researcher@openmined.org\", password=\"****\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a701372-0aca-4379-a50f-62ed76d81b39",
   "metadata": {},
   "source": [
    "## Step 2. Get Mock data and test the code for the Deep Learning (`DL`) experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d327de04-6548-4936-b525-04475a65f42d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mock_data = datasites[\"Cleveland Clinic\"].datasets[\"Heart Disease Dataset\"].assets[\"Heart Study Data\"].mock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02f63383-3182-4cb9-a8fa-87af2c426096",
   "metadata": {},
   "source": [
    "**TYPE ANNOTATIONS**\n",
    "\n",
    "Let's use Python type annotations to better document the input/output of each step in the `DL` experiment.\n",
    "\n",
    "> 💡 All the annotations are **exactly** the same used in the [`ML` experiment](./05-Federated-learning-Experiment.ipynb), but `Dataset` that has been updated to handle data represented as `torch.Tensor` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ce7a79-1d27-410f-bc30-9c98378b0917",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.typing as npt\n",
    "from typing import Union, TypeVar, Any\n",
    "\n",
    "DataFrame = TypeVar(\"pandas.DataFrame\")\n",
    "NDArray = npt.NDArray[Any]\n",
    "NDArrayInt = npt.NDArray[np.int_]\n",
    "NDArrayFloat = npt.NDArray[np.float_]\n",
    "\n",
    "Dataset = TypeVar(\"torch.utils.data.TensorDataset\")  # NEW!\n",
    "Metric = TypeVar(\"Metric\", bound=dict[str, Union[float, NDArrayInt]])\n",
    "Metrics = TypeVar(\"Metrics\", bound=tuple[Metric, Metric])  # train and test\n",
    "ModelParams = TypeVar(\"ModelParams\", bound=dict[str, NDArrayFloat])\n",
    "Result = TypeVar(\"Result\", bound=tuple[Metrics, ModelParams])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e74dea8-dff5-4c68-8ea3-fee065b3512f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dl_experiment(data: DataFrame, model_params: ModelParams = None, training_epochs: int = 10) -> Result:\n",
    "    \"\"\"DL Experiment using a Multi-layer Perceptron (non-linear) Classifier.\n",
    "    Steps:\n",
    "    1. Preprocessing (partitioning; missing values & scaling; convert to Tensor objects)\n",
    "    2. MLP model definition and init (w/ `model_params`)\n",
    "    3. Training loop w/ Loss & Optimizer. Gather updated model parameters.\n",
    "    4. Evaluation: collect metrics on training and test partitions.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    data : DataFrame\n",
    "        Input Heart Study data represented as Pandas DataFrame.\n",
    "    model_params: ModelParams (dict)\n",
    "        DL Model Parameters\n",
    "    training_epochs : int (default = 10)\n",
    "        Number of training epochs\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "        metrics : Metrics\n",
    "            Evaluation metrics (i.e. MCC, Confusion matrix) on both training and test\n",
    "            data partitions.\n",
    "        model_params : ModelParams\n",
    "            Update model params after training (converted as Numpy NDArrays)\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "    import torch as th\n",
    "    \n",
    "    # preprocessing\n",
    "    from sklearn.model_selection import train_test_split\n",
    "    from sklearn.compose import ColumnTransformer\n",
    "    from sklearn.pipeline import Pipeline\n",
    "    from sklearn.impute import SimpleImputer\n",
    "    from sklearn.preprocessing import RobustScaler\n",
    "    from torch.utils.data import TensorDataset\n",
    "    \n",
    "    # training (model) using torch\n",
    "    import torch.nn as nn\n",
    "    import torch.optim as optim\n",
    "    \n",
    "    # evaluation (metrics)\n",
    "    from sklearn.metrics import matthews_corrcoef as mcc\n",
    "    from sklearn.metrics import confusion_matrix\n",
    "\n",
    "    def preprocess(data: DataFrame) -> tuple[Dataset, Dataset]:\n",
    "\n",
    "        def by_demographics(data: DataFrame) -> NDArray:\n",
    "            sex = data[\"sex\"].map(lambda v: \"0\" if v == 0 else \"1\")\n",
    "            target = data[\"num\"].map(lambda v: \"0\" if v == 0 else \"1\")\n",
    "            return (sex + target).values\n",
    "        # Convert all data to float32 arrays for cross torck-kernels compatibility.\n",
    "        X = data.drop(columns=[\"age\", \"sex\", \"num\"], axis=1)\n",
    "        y = (data[\"num\"].map(lambda v: 0 if v == 0 else 1)).to_numpy(dtype=np.float32)\n",
    "    \n",
    "        X_train, X_test, y_train, y_test = train_test_split(\n",
    "            X, y, random_state=12345, stratify=by_demographics(data))\n",
    "        \n",
    "        preprocessor = ColumnTransformer(\n",
    "            transformers=[(\"numerical\",\n",
    "                           Pipeline(steps=[(\"imputer\", SimpleImputer(strategy=\"mean\")),\n",
    "                                           (\"scaler\", RobustScaler()),]),\n",
    "                           [\"trestbps\", \"chol\", \"thalach\", \"oldpeak\"]),\n",
    "                          (\"categorical\", \n",
    "                           SimpleImputer(strategy=\"most_frequent\",),\n",
    "                           [\"ca\", \"cp\", \"exang\", \"fbs\", \"restecg\", \"slope\", \"thal\"])])\n",
    "        \n",
    "        X_train = preprocessor.fit_transform(X_train).astype(np.float32)\n",
    "        X_test  = preprocessor.transform(X_test).astype(np.float32)\n",
    "        # Convert to torch tensor\n",
    "        X_train, X_test = th.from_numpy(X_train), th.from_numpy(X_test)\n",
    "        y_train, y_test = th.from_numpy(y_train), th.from_numpy(y_test)\n",
    "        # reshape tensor to add batch dimension\n",
    "        y_train, y_test = y_train.view(y_train.shape[0], 1), y_test.view(y_test.shape[0], 1)\n",
    "        # gather all data in TensorDataset\n",
    "        return TensorDataset(X_train, y_train), TensorDataset(X_test, y_test)\n",
    "\n",
    "    def train(model: nn.Module, device: th.device, training_data: \"TensorDataset\") -> ModelParams:\n",
    "        train_loader = th.utils.data.DataLoader(training_data, shuffle=True,\n",
    "                                                batch_size=min(200, len(training_data)))\n",
    "        # Loss and optimizer\n",
    "        learning_rate = 0.001\n",
    "        criterion = nn.BCELoss()\n",
    "        optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "        # Train the model\n",
    "        for epoch in range(training_epochs):\n",
    "            for data in train_loader:\n",
    "                X_train, y_train = data[0].to(device), data[1].to(device)\n",
    "                y_pred = model(X_train)\n",
    "                loss = criterion(y_pred, y_train)\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "        # return model parameters as NDArray        \n",
    "        return {k: t.cpu().numpy() for k, t in model.state_dict().items()}\n",
    "\n",
    "    def evaluate(model: nn.Module, device: th.device, dataset: \"TensorDataset\") -> Metric:\n",
    "        data_loader = th.utils.data.DataLoader(dataset, batch_size=min(200, len(dataset)))\n",
    "        with th.no_grad():\n",
    "            y_true, y_pred = [], []\n",
    "            for data in data_loader:\n",
    "                X_val, y_val = data[0].to(device), data[1].to(device)\n",
    "                y_pred.append(model(X_val).round().cpu().detach().numpy())\n",
    "                y_true.append(y_val.cpu().numpy())\n",
    "            y_pred = np.vstack(y_pred)\n",
    "            y_true = np.vstack(y_true)\n",
    "        return {\"mcc\": mcc(y_true, y_pred), \"cm\": confusion_matrix(y_true, y_pred)}\n",
    "    \n",
    "    # -- DL Experiment --\n",
    "    # 1. preprocessing\n",
    "    training_data, test_data = preprocess(data)\n",
    "    # 2. model setup\n",
    "    class MLP(nn.Module):\n",
    "        def __init__(self):\n",
    "            super().__init__()\n",
    "            self.fc1 = nn.Linear(11, 22)\n",
    "            self.fc2 = nn.Linear(22, 11)\n",
    "            self.classifier = nn.Linear(11, 1)\n",
    "    \n",
    "        def forward(self, x):\n",
    "            x = th.sigmoid(self.fc1(x))\n",
    "            x = th.sigmoid(self.fc2(x))\n",
    "            return th.sigmoid(self.classifier(x))\n",
    "\n",
    "    clf = MLP()\n",
    "    if model_params:  # convert to torch tensor and load\n",
    "        clf.load_state_dict({k: th.from_numpy(v) for k, v in model_params.items()})\n",
    "    \n",
    "    device = th.device(\"cuda\" if th.cuda.is_available() else \"mps\" if th.backends.mps.is_available() else \"cpu\")\n",
    "    clf.to(device)\n",
    "    # 3. training\n",
    "    model_params = train(clf, device, training_data)\n",
    "    # 4. evaluation\n",
    "    training_metrics = evaluate(clf, device, training_data)\n",
    "    test_metrics = evaluate(clf, device, test_data)\n",
    "    return (training_metrics, test_metrics), model_params\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0bc84af-273f-4d25-88b4-0e441fbf24b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics, model_params = dl_experiment(data=mock_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eec11f5f-20a7-4107-81c3-108b92eb5647",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bf101b7-6f09-42c3-b634-4862f395f562",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(model_params) == 6  # 3 layers x 2 tensors each (weight + bias)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f46524e-a024-4a04-bf3e-e2261bb19e1c",
   "metadata": {},
   "source": [
    "## Step 3. FL Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "408b7221-1884-4e23-bdf0-42c9a63ca6e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from syft.service.policy.policy import MixedInputPolicy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d38504-f48e-4e27-bbd4-9e10f90ac1c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for name, datasite in datasites.items():\n",
    "    data_asset = datasite.datasets[\"Heart Disease Dataset\"].assets[\"Heart Study Data\"]\n",
    "\n",
    "    syft_fl_experiment = sy.syft_function(\n",
    "        input_policy=MixedInputPolicy(\n",
    "            client=datasite, data=data_asset, model_params=dict, training_epochs=int\n",
    "        )\n",
    "    )(dl_experiment)\n",
    "    dl_training_project = sy.Project(\n",
    "        name=\"DL Experiment for FL\",\n",
    "        description=\"\"\"I would like to run this DL experiment on Heart Disease data, to be used in a Federated Learning fashion. \n",
    "        The function will iteratively be invoked by passing updated model parameters resulting from FL iterations. \n",
    "        These parameters are obtained by averaging the parameters of each model run on each of the four Hospital datasites.\n",
    "        Training and Test metrics will be returned after each execution to monitor the learning progress.\"\"\",\n",
    "        members=[datasite],\n",
    "    )\n",
    "    dl_training_project.create_code_request(syft_fl_experiment, datasite)\n",
    "    project = dl_training_project.send()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd5b6a9a-bae5-4351-afde-6b762f8d57f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import check_status_last_code_requests\n",
    "\n",
    "check_status_last_code_requests(datasites)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de73933-60b5-4602-89b4-11d33de0c74c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "from itertools import groupby\n",
    "from collections import defaultdict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69dfa7d6-43a0-4a77-9b3e-481bd05c57ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def avg(all_model_params: list[ModelParams]) -> ModelParams:\n",
    "    return {layer: np.average([params[layer] for params in all_model_params], axis=0) \n",
    "            for layer in all_model_params[0].keys()}\n",
    "\n",
    "\n",
    "def fl_experiment(datasites, fl_epochs: int = 75, start_epoch: int = 0, \n",
    "                  training_epochs: int = 5, model_params: ModelParams = None):\n",
    "    \"\"\"Federated Learning Experiment, using PyTorch DL models.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    datasites: dict[str, sy.DatasiteClient]\n",
    "        Map of each available Syft Datasite.\n",
    "    fl_epochs: int (default: 75)\n",
    "        Number of FL epochs to run in the experiment\n",
    "    start_epoch: int (default: 0)\n",
    "        Starting FL epoch (useful for consecutive function calls)\n",
    "    training_epochs: int (default: 5)\n",
    "        Number of training epochs of DL models on each datasite (see `dl_experiment`).\n",
    "    model_params: ModelParams (default: None)\n",
    "        If not None, the initial model parameters to consider. This is useful to iteratively\n",
    "        invoke this function multiple times, allowing to store intermediate results.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    fl_metrics: dict[int, list]\n",
    "        Dictionary containing metrics, and model parameters of the DL experiment during each FL epoch.\n",
    "        The list includes one entry per each Datasite.\n",
    "    fl_model_params: ModelParams\n",
    "        The average model parameters resulting from the FL experiment\n",
    "    \"\"\"\n",
    "    fl_model_params = dict() if not model_params else model_params\n",
    "    fl_metrics = defaultdict(list)  # one entry per epoch as a list\n",
    "    for epoch in range(start_epoch, (total_epochs := start_epoch+fl_epochs)):\n",
    "        is_checkpoint = (epoch == start_epoch) or (epoch % (total_epochs // 4) == 0) or (epoch == total_epochs - 1)\n",
    "        for datasite in datasites.values():\n",
    "            data_asset = datasite.datasets[\"Heart Disease Dataset\"].assets[\"Heart Study Data\"]\n",
    "            metrics, params = datasite.code.dl_experiment(data=data_asset, model_params=fl_model_params, \n",
    "                                                          training_epochs=training_epochs).get_from(datasite)\n",
    "            fl_metrics[epoch].append((metrics, params))\n",
    "        if is_checkpoint:\n",
    "            print(\"Epoch: \", epoch)\n",
    "            for idx, name in enumerate(datasites):\n",
    "                metrics = fl_metrics[epoch][idx][0]\n",
    "                print(f\'\\t {name}: Train:{metrics[0][\"mcc\"]:.3f} | Test:{metrics[1][\"mcc\"]:.3f}\')\n",
    "        fl_model_params = avg([params for _, params in fl_metrics[epoch]])\n",
    "    return fl_metrics, fl_model_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef397fa-237f-49b6-9013-289e1a18627f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fl_metrics, fl_model_params = fl_experiment(datasites)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a61324-ed24-43cc-9c01-f8c1372399b9",
   "metadata": {},
   "source": [
    "## Step 4. Plot Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f6e83ea-762f-4426-bfab-c098f8c6d05d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from utils import plot_fl_metrics\n",
    "\n",
    "plot_fl_metrics(datasites, fl_metrics, title=\"FL Experiment on Heart Disease Data w/ MLP\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe612ed5-0549-480e-a23f-79723c681fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import plot_all_confusion_matrices\n",
    "\n",
    "last_epoch = sorted(fl_metrics)[-1]\n",
    "confusion_matrices = {name: fl_metrics[last_epoch][idx][0][1][\"cm\"] for idx, name in enumerate(datasites)}\n",
    "plot_all_confusion_matrices(confusion_matrices, title=\"Confusion Matrices of FL MLP\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e59739-8e38-41f7-bb05-da81d017ba85",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "\n",
    "to complete"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}