src/Standard_CentralNN_DNA+MRI.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import necessary libraries and modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%bash\n",
    "\n",
    "# kaggle datasets download -d nageshsingh/dna-sequence-dataset\n",
    "# mkdir data data/DNA\n",
    "# unzip dna-sequence-dataset.zip -d data/DNA\n",
    "# rm dna-sequence-dataset.zip\n",
    "\n",
    "# kaggle datasets download -d masoudnickparvar/brain-tumor-mri-dataset\n",
    "# mkdir -p data/MRI\n",
    "# unzip brain-tumor-mri-dataset.zip -d data/MRI\n",
    "# rm brain-tumor-mri-dataset.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import time\n",
    "from collections import OrderedDict\n",
    "from typing import (\n",
    "    List, Tuple, Dict, Optional, Callable, Union\n",
    ")\n",
    "\n",
    "import numpy as np\n",
    "import torchvision\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import flwr as fl\n",
    "from flwr.common import (\n",
    "    Metrics, EvaluateIns, EvaluateRes, FitIns, FitRes, MetricsAggregationFn, \n",
    "    Scalar, logger, ndarrays_to_parameters, parameters_to_ndarrays,\n",
    "    Parameters, NDArrays\n",
    ")\n",
    "from flwr.server.client_proxy import ClientProxy\n",
    "from flwr.server.client_manager import ClientManager\n",
    "from flwr.server.strategy.aggregate import aggregate, weighted_loss_avg\n",
    "from logging import WARNING\n",
    "\n",
    "from utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Architecture Creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quantum device definitions for MRI and DNA\n",
    "mri_n_qubits = 4\n",
    "dna_n_qubits = 7\n",
    "expert_vector = (mri_n_qubits+dna_n_qubits) // 2 + 1\n",
    "num_of_expert = 2\n",
    "\n",
    "# Define the MRI network\n",
    "class MRINet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MRINet, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            nn.Conv2d(3, 16, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
    "            nn.Conv2d(16, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        )\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(32 * 56 * 56, 128),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(128, expert_vector),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.classifier(x)\n",
    "        return x\n",
    "\n",
    "class DNANet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(DNANet, self).__init__()        \n",
    "        self.fc1 = nn.Linear(input_sp, 64)\n",
    "        self.fc2 = nn.Linear(64, 32)\n",
    "        self.fc3 = nn.Linear(32, 16)\n",
    "        self.fc4 = nn.Linear(16, 8)\n",
    "        self.fc5 = nn.Linear(8, expert_vector)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Forward pass of the neural network\n",
    "        \"\"\"\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = F.relu(self.fc3(x))\n",
    "        x = F.relu(self.fc4(x))\n",
    "        x = self.fc5(x)\n",
    "        return x\n",
    "\n",
    "class Net(nn.Module):\n",
    "    def __init__(self, num_classes_mri, num_classes_dna):\n",
    "        super(Net, self).__init__()\n",
    "        self.mri_net = MRINet()\n",
    "        self.dna_net = DNANet()\n",
    "        \n",
    "        self.feature_dim = expert_vector\n",
    "        self.num_heads = expert_vector\n",
    "        \n",
    "        self.attention = nn.MultiheadAttention(embed_dim=num_of_expert*self.feature_dim, num_heads=self.num_heads)\n",
    "        self.fc_gate = nn.Linear(num_of_expert*self.feature_dim, 2) \n",
    "        self.fc2_mri = nn.Linear(self.feature_dim, num_classes_mri)\n",
    "        self.fc2_dna = nn.Linear(self.feature_dim, num_classes_dna)\n",
    "        \n",
    "    def forward(self, mri_input, dna_input):\n",
    "        mri_features = self.mri_net(mri_input)\n",
    "        dna_features = self.dna_net(dna_input)\n",
    "        combined_features = torch.cat((mri_features, dna_features), dim=1)\n",
    "        combined_features = combined_features.unsqueeze(0)\n",
    "        attn_output, _ = self.attention(combined_features, combined_features, combined_features)\n",
    "        attn_output = attn_output.squeeze(0)\n",
    "        gates = F.softmax(self.fc_gate(attn_output), dim=1)\n",
    "        combined_output = (gates[:, 0].unsqueeze(1) * mri_features + \n",
    "                           gates[:, 1].unsqueeze(1) * dna_features)\n",
    "        mri_output = self.fc2_mri(combined_output)\n",
    "        dna_output = self.fc2_dna(combined_output)\n",
    "        return mri_output, dna_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the FlowerClient class for federated learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FlowerClient(fl.client.NumPyClient):\n",
    "    def __init__(self, cid, net, trainloader, valloader, device, batch_size, save_results, matrix_path, roc_path,\n",
    "                 yaml_path, classes):\n",
    "        self.net = net\n",
    "        self.trainloader = trainloader\n",
    "        self.valloader = valloader\n",
    "        self.cid = cid\n",
    "        self.device = device\n",
    "        self.batch_size = batch_size\n",
    "        self.save_results = save_results\n",
    "        self.matrix_path = matrix_path\n",
    "        self.roc_path = roc_path\n",
    "        self.yaml_path = yaml_path\n",
    "        self.classes = classes\n",
    "\n",
    "    def get_parameters(self, config):\n",
    "        print(f\"[Client {self.cid}] get_parameters\")\n",
    "        return get_parameters2(self.net)\n",
    "\n",
    "    def fit(self, parameters, config):\n",
    "        server_round = config['server_round']\n",
    "        local_epochs = config['local_epochs']\n",
    "        lr = float(config[\"learning_rate\"])\n",
    "\n",
    "        print(f'[Client {self.cid}, round {server_round}] fit, config: {config}')\n",
    "\n",
    "        set_parameters(self.net, parameters)\n",
    "\n",
    "        criterion_mri = torch.nn.CrossEntropyLoss()\n",
    "        criterion_dna = torch.nn.CrossEntropyLoss()\n",
    "        optimizer = torch.optim.Adam(self.net.parameters(), lr=lr)\n",
    "\n",
    "        results = engine.train(self.net, self.trainloader, self.valloader, optimizer=optimizer, loss_fn=(criterion_mri, criterion_dna),\n",
    "                               epochs=local_epochs, device=self.device, task=\"Multimodal\")\n",
    "\n",
    "        if self.save_results:\n",
    "            save_graphs_multimodal(self.save_results, local_epochs, results, f\"_Client {self.cid}\")\n",
    "\n",
    "        return get_parameters2(self.net), len(self.trainloader), {}\n",
    "\n",
    "    def evaluate(self, parameters, config):\n",
    "        print(f\"[Client {self.cid}] evaluate, config: {config}\")\n",
    "        set_parameters(self.net, parameters)\n",
    "\n",
    "        loss, accuracy, y_pred, y_true, y_proba = engine.test_multimodal_health(self.net, self.valloader,\n",
    "                                                              loss_fn=(torch.nn.CrossEntropyLoss(),torch.nn.CrossEntropyLoss()), device=self.device)\n",
    "\n",
    "        loss_mri, loss_dna = loss\n",
    "        accuracy_mri, accuracy_dna = accuracy\n",
    "        y_pred_mri, y_pred_dna = y_pred\n",
    "        y_true_mri, y_true_dna = y_true\n",
    "        y_proba_mri, y_proba_dna = y_proba\n",
    "\n",
    "        if self.save_results:\n",
    "            os.makedirs(self.save_results, exist_ok=True)\n",
    "            if self.matrix_path:\n",
    "                save_matrix(y_true_mri, y_pred_mri, self.save_results + \"MRI_\" + self.matrix_path, self.classes[0])\n",
    "                save_matrix(y_true_dna, y_pred_dna, self.save_results + \"DNA_\" + self.matrix_path, self.classes[1])\n",
    "            if self.roc_path:\n",
    "                save_roc(y_true_mri, y_proba_mri, self.save_results + \"MRI_\" + self.roc_path, len(self.classes[0]))\n",
    "                save_roc(y_true_dna, y_proba_dna, self.save_results + \"DNA_\" + self.roc_path, len(self.classes[1]))\n",
    "\n",
    "        return float(loss_mri), len(self.valloader), {\"accuracy\": float(accuracy_mri)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the client_common function to set up the Flower client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def client_common(cid, model_save, path_yaml, path_roc, results_save, path_matrix,\n",
    "                  batch_size, trainloaders, valloaders, DEVICE, CLASSES):\n",
    "    trainloader = trainloaders[int(cid)]\n",
    "    valloader = valloaders[int(cid)]\n",
    "\n",
    "    net = Net(len(CLASSES[0]), len(CLASSES[1])).to(DEVICE)\n",
    "\n",
    "    if os.path.exists(model_save):\n",
    "        print(\" To get the checkpoint\")\n",
    "        checkpoint = torch.load(model_save, map_location=DEVICE)['model_state_dict']\n",
    "        net.load_state_dict(checkpoint)\n",
    "\n",
    "    return FlowerClient(cid, net, trainloader, valloader, device=DEVICE, batch_size=batch_size,\n",
    "                        matrix_path=path_matrix, roc_path=path_roc, save_results=results_save, yaml_path=path_yaml,\n",
    "                        classes=CLASSES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define utility functions for federated learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:\n",
    "    accuracies = [num_examples * m[\"accuracy\"] for num_examples, m in metrics]\n",
    "    examples = [num_examples for num_examples, _ in metrics]\n",
    "    return {\"accuracy\": sum(accuracies) / sum(examples)}\n",
    "\n",
    "def evaluate2(server_round: int, parameters: NDArrays,\n",
    "              config: Dict[str, Scalar]) -> Optional[Tuple[float, Dict[str, Scalar]]]:\n",
    "    set_parameters(central, parameters)\n",
    "    loss, accuracy, y_pred, y_true, y_proba = engine.test_multimodal_health(central, testloader, loss_fn=(torch.nn.CrossEntropyLoss(),torch.nn.CrossEntropyLoss()),\n",
    "                                                          device=DEVICE)\n",
    "    loss_mri, loss_dna = loss\n",
    "    accuracy_mri, accuracy_dna = accuracy\n",
    "    y_pred_mri, y_pred_dna = y_pred\n",
    "    y_true_mri, y_true_dna = y_true\n",
    "    y_proba_mri, y_proba_dna = y_proba\n",
    "    os.makedirs(save_results, exist_ok=True)\n",
    "    save_matrix(y_true_mri, y_pred_mri, save_results + \"mri_confusion_matrix_test.png\", CLASSES[0])\n",
    "    save_roc(y_true_mri, y_proba_mri, save_results + \"mri_roc_test.png\", len(CLASSES[0]))\n",
    "    save_matrix(y_true_dna, y_pred_dna, save_results + \"dna_confusion_matrix_test.png\", CLASSES[1])\n",
    "    save_roc(y_true_dna, y_proba_dna, save_results + \"dna_roc_test.png\", len(CLASSES[1]))\n",
    "\n",
    "    print(f\"Server-side evaluation MRI loss {loss_mri} / MRI accuracy {accuracy_mri}\")\n",
    "    print(f\"Server-side evaluation DNA loss {loss_dna} / DNA accuracy {accuracy_dna}\")\n",
    "    return (loss_mri, loss_dna), {\"accuracy\": (accuracy_mri, accuracy_dna)}\n",
    "\n",
    "def get_on_fit_config_fn(epoch=2, lr=0.001, batch_size=32) -> Callable[[int], Dict[str, str]]:\n",
    "    def fit_config(server_round: int) -> Dict[str, str]:\n",
    "        config = {\n",
    "            \"learning_rate\": str(lr),\n",
    "            \"batch_size\": str(batch_size),\n",
    "            \"server_round\": server_round,\n",
    "            \"local_epochs\": epoch\n",
    "        }\n",
    "        return config\n",
    "    return fit_config\n",
    "\n",
    "def aggreg_fit_checkpoint(server_round, aggregated_parameters, central_model, path_checkpoint):\n",
    "    if aggregated_parameters is not None:\n",
    "        print(f\"Saving round {server_round} aggregated_parameters...\")\n",
    "        aggregated_ndarrays: List[np.ndarray] = parameters_to_ndarrays(aggregated_parameters)        \n",
    "        params_dict = zip(central_model.state_dict().keys(), aggregated_ndarrays)\n",
    "        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})\n",
    "        central_model.load_state_dict(state_dict, strict=True)\n",
    "        if path_checkpoint:\n",
    "            torch.save({\n",
    "                'model_state_dict': central_model.state_dict(),\n",
    "            }, path_checkpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the FedCustom strategy class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A Strategy from scratch with the same sampling of the clients as it is in FedAvg\n",
    "# and then change the configuration dictionary\n",
    "class FedCustom(fl.server.strategy.Strategy):\n",
    "    def __init__(\n",
    "            self,\n",
    "            fraction_fit: float = 1.0,\n",
    "            fraction_evaluate: float = 1.0,\n",
    "            min_fit_clients: int = 2,\n",
    "            min_evaluate_clients: int = 2,\n",
    "            min_available_clients: int = 2,\n",
    "            evaluate_fn: Optional[\n",
    "                    Callable[[int, NDArrays, Dict[str, Scalar]], Optional[Tuple[float, Dict[str, Scalar]]]]\n",
    "                ] = None,\n",
    "            on_fit_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,\n",
    "            on_evaluate_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,\n",
    "            accept_failures: bool = True,\n",
    "            initial_parameters: Optional[Parameters] = None,\n",
    "            fit_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,\n",
    "            evaluate_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,\n",
    "    ) -> None:\n",
    "        super().__init__()\n",
    "        self.fraction_fit = fraction_fit\n",
    "        self.fraction_evaluate = fraction_evaluate\n",
    "        self.min_fit_clients = min_fit_clients\n",
    "        self.min_evaluate_clients = min_evaluate_clients\n",
    "        self.min_available_clients = min_available_clients\n",
    "        self.evaluate_fn = evaluate_fn\n",
    "        self.on_fit_config_fn = on_fit_config_fn\n",
    "        self.on_evaluate_config_fn = on_evaluate_config_fn,\n",
    "        self.accept_failures = accept_failures\n",
    "        self.initial_parameters = initial_parameters\n",
    "        self.fit_metrics_aggregation_fn = fit_metrics_aggregation_fn\n",
    "        self.evaluate_metrics_aggregation_fn = evaluate_metrics_aggregation_fn\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        return f\"FedCustom (accept_failures={self.accept_failures})\"\n",
    "\n",
    "    def initialize_parameters(\n",
    "        self, client_manager: ClientManager\n",
    "    ) -> Optional[Parameters]:\n",
    "        \"\"\"Initialize global model parameters.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        initial_parameters = self.initial_parameters\n",
    "        self.initial_parameters = None  # Don't keep initial parameters in memory\n",
    "        return initial_parameters\n",
    "\n",
    "    def num_fit_clients(self, num_available_clients: int) -> Tuple[int, int]:\n",
    "        \"\"\"Return sample size and required number of clients.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        num_clients = int(num_available_clients * self.fraction_fit)\n",
    "        return max(num_clients, self.min_fit_clients), self.min_available_clients\n",
    "\n",
    "    def configure_fit(\n",
    "        self, server_round: int, parameters: Parameters, client_manager: ClientManager\n",
    "    ) -> List[Tuple[ClientProxy, FitIns]]:\n",
    "        \"\"\"Configure the next round of training.\"\"\"\n",
    "        # Sample clients\n",
    "        sample_size, min_num_clients = self.num_fit_clients(\n",
    "            client_manager.num_available()\n",
    "        )\n",
    "\n",
    "        clients = client_manager.sample(\n",
    "            num_clients=sample_size, min_num_clients=min_num_clients\n",
    "        )\n",
    "        # Create custom configs\n",
    "        n_clients = len(clients)\n",
    "        half_clients = n_clients // 2\n",
    "        # Custom fit config function provided\n",
    "        standard_lr = lr\n",
    "        higher_lr = 0.003\n",
    "        config = {\"server_round\": server_round, \"local_epochs\": 1}\n",
    "        if self.on_fit_config_fn is not None:\n",
    "            # Custom fit config function provided\n",
    "            config = self.on_fit_config_fn(server_round)\n",
    "\n",
    "        # fit_ins = FitIns(parameters, config)\n",
    "        # Return client/config pairs\n",
    "        fit_configurations = []\n",
    "        for idx, client in enumerate(clients):\n",
    "            config[\"learning_rate\"] = standard_lr if idx < half_clients else higher_lr\n",
    "            \"\"\"\n",
    "            Each pair of (ClientProxy, FitRes) constitutes \n",
    "            a successful update from one of the previously selected clients.\n",
    "            \"\"\"\n",
    "            fit_configurations.append(\n",
    "                (\n",
    "                    client,\n",
    "                    FitIns(\n",
    "                        parameters,\n",
    "                        config\n",
    "                    )\n",
    "                )\n",
    "            )\n",
    "        # Successful updates from the previously selected and configured clients\n",
    "        return fit_configurations\n",
    "\n",
    "    def aggregate_fit(\n",
    "        self,\n",
    "        server_round: int,\n",
    "        results: List[Tuple[ClientProxy, FitRes]],\n",
    "        failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]],\n",
    "    ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:\n",
    "        \"\"\"Aggregate fit results using weighted average. (each round)\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        if not results:\n",
    "            return None, {}\n",
    "\n",
    "        # Do not aggregate if there are failures and failures are not accepted\n",
    "        if not self.accept_failures and failures:\n",
    "            return None, {}\n",
    "\n",
    "        # Convert results parameters --> array matrix\n",
    "        weights_results = [\n",
    "            (parameters_to_ndarrays(fit_res.parameters), fit_res.num_examples)\n",
    "            for _, fit_res in results\n",
    "        ]\n",
    "\n",
    "        # Aggregate parameters using weighted average between the clients and convert back to parameters object (bytes)\n",
    "        parameters_aggregated = ndarrays_to_parameters(aggregate(weights_results))\n",
    "\n",
    "        metrics_aggregated = {}\n",
    "        # Aggregate custom metrics if aggregation fn was provided\n",
    "        if self.fit_metrics_aggregation_fn:\n",
    "            fit_metrics = [(res.num_examples, res.metrics) for _, res in results]\n",
    "            metrics_aggregated = self.fit_metrics_aggregation_fn(fit_metrics)\n",
    "\n",
    "        elif server_round == 1:  # Only log this warning once\n",
    "            logger.log(WARNING, \"No fit_metrics_aggregation_fn provided\")\n",
    "\n",
    "        # Same function as SaveModelStrategy(fl.server.strategy.FedAvg)\n",
    "        \"\"\"Aggregate model weights using weighted average and store checkpoint\"\"\"\n",
    "        aggreg_fit_checkpoint(server_round, parameters_aggregated, central, model_save)\n",
    "        return parameters_aggregated, metrics_aggregated\n",
    "\n",
    "    def num_evaluation_clients(self, num_available_clients: int) -> Tuple[int, int]:\n",
    "        \"\"\"Use a fraction of available clients for evaluation.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        num_clients = int(num_available_clients * self.fraction_evaluate)\n",
    "        return max(num_clients, self.min_evaluate_clients), self.min_available_clients\n",
    "\n",
    "    def configure_evaluate(\n",
    "        self, server_round: int, parameters: Parameters, client_manager: ClientManager\n",
    "    ) -> List[Tuple[ClientProxy, EvaluateIns]]:\n",
    "        \"\"\"Configure the next round of evaluation.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        # Do not configure federated evaluation if fraction eval is 0.\n",
    "        if self.fraction_evaluate == 0.0:\n",
    "            return []\n",
    "\n",
    "        # Parameters and config\n",
    "        config = {}  # {\"server_round\": server_round, \"local_epochs\": 1}\n",
    "\n",
    "        evaluate_ins = EvaluateIns(parameters, config)\n",
    "\n",
    "        # Sample clients\n",
    "        sample_size, min_num_clients = self.num_evaluation_clients(\n",
    "            client_manager.num_available()\n",
    "        )\n",
    "\n",
    "        clients = client_manager.sample(\n",
    "            num_clients=sample_size, min_num_clients=min_num_clients\n",
    "        )\n",
    "\n",
    "        # Return client/config pairs\n",
    "        # Each pair of (ClientProxy, FitRes) constitutes a successful update from one of the previously selected clients\n",
    "        return [(client, evaluate_ins) for client in clients]\n",
    "\n",
    "    def aggregate_evaluate(\n",
    "        self,\n",
    "        server_round: int,\n",
    "        results: List[Tuple[ClientProxy, EvaluateRes]],\n",
    "        failures: List[Union[Tuple[ClientProxy, EvaluateRes], BaseException]],\n",
    "    ) -> Tuple[Optional[float], Dict[str, Scalar]]:\n",
    "        \"\"\"Aggregate evaluation losses using weighted average.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        if not results:\n",
    "            return None, {}\n",
    "\n",
    "        # Do not aggregate if there are failures and failures are not accepted\n",
    "        if not self.accept_failures and failures:\n",
    "            return None, {}\n",
    "\n",
    "        # Aggregate loss\n",
    "        loss_aggregated = weighted_loss_avg(\n",
    "            [\n",
    "                (evaluate_res.num_examples, evaluate_res.loss)\n",
    "                for _, evaluate_res in results\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        metrics_aggregated = {}\n",
    "        # Aggregate custom metrics if aggregation fn was provided\n",
    "        if self.evaluate_metrics_aggregation_fn:\n",
    "            eval_metrics = [(res.num_examples, res.metrics) for _, res in results]\n",
    "            metrics_aggregated = self.evaluate_metrics_aggregation_fn(eval_metrics)\n",
    "\n",
    "        # Only log this warning once\n",
    "        elif server_round == 1:\n",
    "            logger.log(WARNING, \"No evaluate_metrics_aggregation_fn provided\")\n",
    "\n",
    "        return loss_aggregated, metrics_aggregated\n",
    "\n",
    "    def evaluate(\n",
    "        self, server_round: int, parameters: Parameters\n",
    "    ) -> Optional[Tuple[float, Dict[str, Scalar]]]:\n",
    "        \"\"\"Evaluate global model parameters using an evaluation function.\"\"\"\n",
    "        # Same function as FedAvg(Strategy)\n",
    "        if self.evaluate_fn is None:\n",
    "            # Let's assume we won't perform the global model evaluation on the server side.\n",
    "            return None\n",
    "\n",
    "        # if we have a global model evaluation on the server side :\n",
    "        parameters_ndarrays = parameters_to_ndarrays(parameters)\n",
    "        eval_res = self.evaluate_fn(server_round, parameters_ndarrays, {})\n",
    "\n",
    "        # if you haven't results\n",
    "        if eval_res is None:\n",
    "            return None\n",
    "\n",
    "        loss, metrics = eval_res\n",
    "        return loss, metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the federated learning strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up your variables directly\n",
    "data_path = 'data/'\n",
    "dataset = 'DNA+MRI'\n",
    "yaml_path = './results/CL_DNA+MRI/results.yml'\n",
    "seed = 0\n",
    "num_workers = 0\n",
    "max_epochs = 25\n",
    "batch_size = 32\n",
    "splitter = 10\n",
    "device = 'gpu'\n",
    "number_clients = 1\n",
    "save_results = 'results/CL_DNA+MRI/'\n",
    "matrix_path = 'confusion_matrix.png'\n",
    "roc_path = 'roc.png'\n",
    "model_save = 'results/CL_DNA+MRI/DNA+MRI_cl.pt'\n",
    "min_fit_clients = 1\n",
    "min_avail_clients = 1\n",
    "min_eval_clients = 1\n",
    "rounds = 1\n",
    "frac_fit = 1.0\n",
    "frac_eval = 0.5\n",
    "lr = 1e-3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DNA+MRI\n",
      "The training set is created for the classes: \n",
      "('glioma', 'meningioma', 'notumor', 'pituitary')\n",
      "('0', '1', '2', '3', '4', '5', '6')\n"
     ]
    }
   ],
   "source": [
    "DEVICE = torch.device(choice_device(device))\n",
    "CLASSES = classes_string(dataset)\n",
    "trainloaders, valloaders, testloader = data_setup.load_datasets(num_clients=number_clients, batch_size=batch_size, resize=224, seed=seed, num_workers=num_workers, splitter=splitter, dataset=dataset, data_path=data_path, data_path_val=None)\n",
    "_, input_sp = next(iter(testloader))[0][1].shape\n",
    "central = Net(len(CLASSES[0]), len(CLASSES[1])).to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy = FedCustom(\n",
    "    fraction_fit=frac_fit,\n",
    "    fraction_evaluate=frac_eval,\n",
    "    min_fit_clients=min_fit_clients,\n",
    "    min_evaluate_clients=min_eval_clients if min_eval_clients else number_clients // 2,\n",
    "    min_available_clients=min_avail_clients,\n",
    "    evaluate_metrics_aggregation_fn=weighted_average,\n",
    "    initial_parameters=ndarrays_to_parameters(get_parameters2(central)),\n",
    "    evaluate_fn=evaluate2,\n",
    "    on_fit_config_fn=get_on_fit_config_fn(epoch=max_epochs, batch_size=batch_size),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DNA+MRI\n",
      "The training set is created for the classes: \n",
      "('glioma', 'meningioma', 'notumor', 'pituitary')\n",
      "('0', '1', '2', '3', '4', '5', '6')\n"
     ]
    }
   ],
   "source": [
    "trainloaders, valloaders, testloader = data_setup.load_datasets(num_clients=number_clients,\n",
    "                                                                batch_size=batch_size,\n",
    "                                                                resize=224,\n",
    "                                                                seed=seed,\n",
    "                                                                num_workers=num_workers,\n",
    "                                                                splitter=splitter,\n",
    "                                                                dataset=dataset,  # Use the specified dataset\n",
    "                                                                data_path=data_path,\n",
    "                                                                data_path_val=None)  # Use the same path for validation data\n",
    "\n",
    "def client_fn(cid: str) -> FlowerClient:\n",
    "    return client_common(cid,\n",
    "                         model_save, path_yaml, path_roc, results_save, path_matrix,\n",
    "                         batch_size, trainloaders, valloaders, DEVICE, CLASSES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the client_fn function and set up the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO flwr 2024-10-06 15:17:35,356 | app.py:175 | Starting Flower simulation, config: ServerConfig(num_rounds=1, round_timeout=None)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "flwr 1.5.0\n",
      "numpy 1.26.4\n",
      "torch 2.4.0+cu121\n",
      "torchvision 0.19.0+cu121\n",
      "Training on cuda:0\n",
      "Start simulation\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-10-06 15:17:37,447\tWARNING utils.py:592 -- Ray currently does not support initializing Ray with fractional cpus. Your num_cpus will be truncated from 11.5 to 11.\n",
      "2024-10-06 15:17:38,940\tINFO worker.py:1786 -- Started a local Ray instance.\n",
      "INFO flwr 2024-10-06 15:17:39,519 | app.py:210 | Flower VCE: Ray initialized with resources: {'accelerator_type:A100': 1.0, 'memory': 125622668698.0, 'node:__internal_head__': 1.0, 'object_store_memory': 58124000870.0, 'CPU': 11.0, 'node:10.42.22.191': 1.0, 'GPU': 1.0}\n",
      "INFO flwr 2024-10-06 15:17:39,520 | app.py:224 | Flower VCE: Resources for each Virtual Client: {'num_cpus': 1, 'num_gpus': 1}\n",
      "INFO flwr 2024-10-06 15:17:39,532 | app.py:270 | Flower VCE: Creating VirtualClientEngineActorPool with 1 actors\n",
      "INFO flwr 2024-10-06 15:17:39,532 | server.py:89 | Initializing global parameters\n",
      "INFO flwr 2024-10-06 15:17:39,533 | server.py:272 | Using initial parameters provided by strategy\n",
      "INFO flwr 2024-10-06 15:17:39,533 | server.py:91 | Evaluating initial parameters\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Updated model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO flwr 2024-10-06 15:17:56,149 | server.py:94 | initial parameters (loss, other metrics): (1.398712741477149, 1.9372813403606415), {'accuracy': (24.479166666666668, 27.269345238095237)}\n",
      "INFO flwr 2024-10-06 15:17:56,150 | server.py:104 | FL starting\n",
      "DEBUG flwr 2024-10-06 15:17:56,150 | server.py:222 | fit_round 1: strategy sampled 1 clients (out of 1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server-side evaluation MRI loss 1.398712741477149 / MRI accuracy 24.479166666666668\n",
      "Server-side evaluation DNA loss 1.9372813403606415 / DNA accuracy 27.269345238095237\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m /home/cudaq/.local/lib/python3.10/site-packages/torch/storage.py:414: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m   return torch.load(io.BytesIO(b))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m [Client 0, round 1] fit, config: {'learning_rate': 0.003, 'batch_size': '32', 'server_round': 1, 'local_epochs': 25}\n",
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m Updated model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|\u001b[34m          \u001b[0m| 0/25 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 1 \tTrain_loss_mri: 1.3339 | Train_acc_mri: 53.1601 | Train_loss_dna: 2.2033 | Train_acc_dna: 32.4144 | Validation_loss_mri: 0.7872 | Validation_acc_mri: 66.8561 | Validation_loss_dna: 1.6118 | Validation_acc_dna: 36.5530\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|\u001b[34m▍         \u001b[0m| 1/25 [00:15<06:06, 15.26s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 2 \tTrain_loss_mri: 0.5563 | Train_acc_mri: 79.5981 | Train_loss_dna: 1.2033 | Train_acc_dna: 53.9703 | Validation_loss_mri: 0.4810 | Validation_acc_mri: 80.8712 | Validation_loss_dna: 1.1740 | Validation_acc_dna: 51.1742\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|\u001b[34m▊         \u001b[0m| 2/25 [00:29<05:31, 14.41s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 3 \tTrain_loss_mri: 0.2988 | Train_acc_mri: 89.2186 | Train_loss_dna: 0.5180 | Train_acc_dna: 83.1264 | Validation_loss_mri: 0.7819 | Validation_acc_mri: 74.2235 | Validation_loss_dna: 0.3835 | Validation_acc_dna: 88.5795\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|\u001b[34m█▏        \u001b[0m| 3/25 [00:55<07:20, 20.01s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 4 \tTrain_loss_mri: 0.1892 | Train_acc_mri: 93.5676 | Train_loss_dna: 0.1699 | Train_acc_dna: 95.8965 | Validation_loss_mri: 0.4628 | Validation_acc_mri: 85.7386 | Validation_loss_dna: 0.2688 | Validation_acc_dna: 92.5568\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|\u001b[34m█▌        \u001b[0m| 4/25 [01:10<06:13, 17.77s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 5 \tTrain_loss_mri: 0.0765 | Train_acc_mri: 97.3801 | Train_loss_dna: 0.0340 | Train_acc_dna: 99.3056 | Validation_loss_mri: 0.5564 | Validation_acc_mri: 85.9848 | Validation_loss_dna: 0.2444 | Validation_acc_dna: 93.1439\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|\u001b[34m██        \u001b[0m| 5/25 [01:26<05:48, 17.42s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 6 \tTrain_loss_mri: 0.0504 | Train_acc_mri: 98.4217 | Train_loss_dna: 0.0168 | Train_acc_dna: 99.7790 | Validation_loss_mri: 0.6312 | Validation_acc_mri: 85.7765 | Validation_loss_dna: 0.2379 | Validation_acc_dna: 93.1250\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|\u001b[34m██▍       \u001b[0m| 6/25 [01:44<05:34, 17.58s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 7 \tTrain_loss_mri: 0.0493 | Train_acc_mri: 98.3025 | Train_loss_dna: 0.0107 | Train_acc_dna: 99.8106 | Validation_loss_mri: 0.7709 | Validation_acc_mri: 84.3371 | Validation_loss_dna: 0.3190 | Validation_acc_dna: 91.4583\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|\u001b[34m██▊       \u001b[0m| 7/25 [01:59<05:01, 16.75s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 8 \tTrain_loss_mri: 0.0329 | Train_acc_mri: 99.0215 | Train_loss_dna: 0.0309 | Train_acc_dna: 99.3687 | Validation_loss_mri: 0.7554 | Validation_acc_mri: 86.0227 | Validation_loss_dna: 0.2589 | Validation_acc_dna: 93.1250\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|\u001b[34m███▏      \u001b[0m| 8/25 [02:13<04:27, 15.72s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 9 \tTrain_loss_mri: 0.0213 | Train_acc_mri: 99.5896 | Train_loss_dna: 0.0209 | Train_acc_dna: 99.5265 | Validation_loss_mri: 0.6754 | Validation_acc_mri: 85.1515 | Validation_loss_dna: 0.3010 | Validation_acc_dna: 89.9811\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 36%|\u001b[34m███▌      \u001b[0m| 9/25 [02:26<04:00, 15.02s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 10 \tTrain_loss_mri: 0.0152 | Train_acc_mri: 99.7159 | Train_loss_dna: 0.0197 | Train_acc_dna: 99.5265 | Validation_loss_mri: 0.6931 | Validation_acc_mri: 87.7273 | Validation_loss_dna: 0.2621 | Validation_acc_dna: 92.5568\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|\u001b[34m████      \u001b[0m| 10/25 [02:39<03:36, 14.45s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 11 \tTrain_loss_mri: 0.0020 | Train_acc_mri: 99.9684 | Train_loss_dna: 0.0013 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.8478 | Validation_acc_mri: 88.6174 | Validation_loss_dna: 0.2368 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|\u001b[34m████▍     \u001b[0m| 11/25 [02:53<03:17, 14.11s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 12 \tTrain_loss_mri: 0.0004 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0005 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.8748 | Validation_acc_mri: 88.9015 | Validation_loss_dna: 0.2545 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 48%|\u001b[34m████▊     \u001b[0m| 12/25 [03:06<02:59, 13.80s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 13 \tTrain_loss_mri: 0.0002 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0003 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.8994 | Validation_acc_mri: 88.9015 | Validation_loss_dna: 0.2578 | Validation_acc_dna: 93.9773\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 52%|\u001b[34m█████▏    \u001b[0m| 13/25 [03:20<02:44, 13.74s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 14 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0002 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9219 | Validation_acc_mri: 88.9015 | Validation_loss_dna: 0.2686 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 56%|\u001b[34m█████▌    \u001b[0m| 14/25 [03:34<02:35, 14.10s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 15 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0002 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9349 | Validation_acc_mri: 88.9015 | Validation_loss_dna: 0.2679 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|\u001b[34m██████    \u001b[0m| 15/25 [03:47<02:16, 13.69s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 16 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9494 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2730 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 64%|\u001b[34m██████▍   \u001b[0m| 16/25 [04:01<02:02, 13.64s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 17 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9629 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2784 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|\u001b[34m██████▊   \u001b[0m| 17/25 [04:14<01:47, 13.47s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 18 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9748 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2812 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 72%|\u001b[34m███████▏  \u001b[0m| 18/25 [04:27<01:33, 13.37s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 19 \tTrain_loss_mri: 0.0001 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9853 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2833 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 76%|\u001b[34m███████▌  \u001b[0m| 19/25 [04:40<01:19, 13.27s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 20 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 0.9998 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2871 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 80%|\u001b[34m████████  \u001b[0m| 20/25 [04:53<01:06, 13.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 21 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 1.0137 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2926 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 84%|\u001b[34m████████▍ \u001b[0m| 21/25 [05:06<00:52, 13.20s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 22 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 1.0249 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2944 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 88%|\u001b[34m████████▊ \u001b[0m| 22/25 [05:20<00:39, 13.25s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 23 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0001 | Train_acc_dna: 100.0000 | Validation_loss_mri: 1.0328 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2948 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|\u001b[34m█████████▏\u001b[0m| 23/25 [05:33<00:26, 13.30s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 24 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0000 | Train_acc_dna: 100.0000 | Validation_loss_mri: 1.0428 | Validation_acc_mri: 88.5985 | Validation_loss_dna: 0.2988 | Validation_acc_dna: 94.5644\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 96%|\u001b[34m█████████▌\u001b[0m| 24/25 [05:51<00:14, 14.57s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m \tTrain Epoch: 25 \tTrain_loss_mri: 0.0000 | Train_acc_mri: 100.0000 | Train_loss_dna: 0.0000 | Train_acc_dna: 100.0000 | Validation_loss_mri: 1.0491 | Validation_acc_mri: 88.3144 | Validation_loss_dna: 0.2971 | Validation_acc_dna: 94.5644\n",
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m save graph in  results/CL_DNA+MRI/\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|\u001b[34m██████████\u001b[0m| 25/25 [06:14<00:00, 14.98s/it]\n",
      "DEBUG flwr 2024-10-06 15:24:38,072 | server.py:236 | fit_round 1 received 1 results and 0 failures\n",
      "WARNING flwr 2024-10-06 15:24:38,241 | 1650717286.py:129 | No fit_metrics_aggregation_fn provided\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving round 1 aggregated_parameters...\n",
      "Updated model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO flwr 2024-10-06 15:25:18,883 | server.py:125 | fit progress: (1, (0.8294942198055131, 0.4351768440433911), {'accuracy': (89.76934523809524, 91.55505952380952)}, 442.73297652590554)\n",
      "DEBUG flwr 2024-10-06 15:25:18,885 | server.py:173 | evaluate_round 1: strategy sampled 1 clients (out of 1)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Server-side evaluation MRI loss 0.8294942198055131 / MRI accuracy 89.76934523809524\n",
      "Server-side evaluation DNA loss 0.4351768440433911 / DNA accuracy 91.55505952380952\n",
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m  To get the checkpoint\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m /tmp/ipykernel_27867/752934157.py:10: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m [Client 0] evaluate, config: {}\n",
      "\u001b[36m(DefaultActor pid=28500)\u001b[0m Updated model\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "DEBUG flwr 2024-10-06 15:25:46,013 | server.py:187 | evaluate_round 1 received 1 results and 0 failures\n",
      "INFO flwr 2024-10-06 15:25:46,014 | server.py:153 | FL finished in 469.8639951709192\n",
      "INFO flwr 2024-10-06 15:25:46,017 | app.py:225 | app_fit: losses_distributed [(1, 1.0490737069736829)]\n",
      "INFO flwr 2024-10-06 15:25:46,017 | app.py:226 | app_fit: metrics_distributed_fit {}\n",
      "INFO flwr 2024-10-06 15:25:46,017 | app.py:227 | app_fit: metrics_distributed {'accuracy': [(1, 88.31439393939394)]}\n",
      "INFO flwr 2024-10-06 15:25:46,018 | app.py:228 | app_fit: losses_centralized [(0, (1.398712741477149, 1.9372813403606415)), (1, (0.8294942198055131, 0.4351768440433911))]\n",
      "INFO flwr 2024-10-06 15:25:46,018 | app.py:229 | app_fit: metrics_centralized {'accuracy': [(0, (24.479166666666668, 27.269345238095237)), (1, (89.76934523809524, 91.55505952380952))]}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Simulation Time = 490.67330741882324 seconds\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.simplefilter(\"ignore\")\n",
    "\n",
    "print(\"flwr\", fl.__version__)\n",
    "print(\"numpy\", np.__version__)\n",
    "print(\"torch\", torch.__version__)\n",
    "print(\"torchvision\", torchvision.__version__)\n",
    "print(f\"Training on {DEVICE}\")\n",
    "\n",
    "client_resources = {\"num_cpus\": 1}\n",
    "\n",
    "if DEVICE.type == \"cuda\":\n",
    "    client_resources[\"num_gpus\"] = 1\n",
    "\n",
    "model_save = model_save\n",
    "path_yaml = yaml_path\n",
    "path_roc = roc_path\n",
    "results_save = save_results\n",
    "path_matrix = matrix_path\n",
    "batch_size = batch_size\n",
    "\n",
    "print(\"Start simulation\")\n",
    "start_simulation = time.time()\n",
    "fl.simulation.start_simulation(\n",
    "    client_fn=client_fn,\n",
    "    num_clients=number_clients,\n",
    "    config=fl.server.ServerConfig(num_rounds=rounds),\n",
    "    strategy=strategy,\n",
    "    client_resources=client_resources\n",
    ")\n",
    "print(f\"Simulation Time = {time.time() - start_simulation} seconds\")"
   ]
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}