MEDfl 2.0.5.dev8__tar.gz → 2.0.6__tar.gz

{medfl-2.0.5.dev8 → medfl-2.0.6}/MEDfl/LearningManager/client.py

@@ -6,6 +6,7 @@ from torch.utils.data import DataLoader
 from .model import Model
 from .utils import params
 import torch
+import numpy as np
 
 class FlowerClient(fl.client.NumPyClient):
     """
@@ -52,6 +53,9 @@ class FlowerClient(fl.client.NumPyClient):
         self.epsilons = []
         self.accuracies = []
         self.losses = []
+        self.train_round_losses = []
+        self.train_round_metrics = []
+
         if self.diff_priv:
             model, optimizer, self.trainloader = self.privacy_engine.make_private_with_epsilon(
                 module=self.local_model.model.train(),
@@ -112,8 +116,11 @@ class FlowerClient(fl.client.NumPyClient):
         # print(config['epochs'])
         print('\n -------------------------------- \n  ')
         self.local_model.set_parameters(parameters)
+
+        epoch_losses, epoch_metrics = [], []
+
         for _ in range(params["train_epochs"]):
-            epsilon = self.local_model.train(
+            epsilon, tr_loss, tr_metric = self.local_model.train(
                 self.trainloader,
                 epoch=_,
                 device=self.device,
@@ -121,30 +128,43 @@ class FlowerClient(fl.client.NumPyClient):
                 diff_priv=self.diff_priv,
             )
             self.epsilons.append(epsilon)
+            epoch_losses.append(tr_loss)
+            epoch_metrics.append(tr_metric)
+
         print(f"epsilon of client {self.cid} : eps = {epsilon}")
+        round_train_loss = float(np.nanmean(epoch_losses)) if len(epoch_losses) else float("nan") 
+         
+        round_train_metric = float(np.nanmean(epoch_metrics)) if len(epoch_metrics) else float("nan")
+
+        self.train_round_losses.append(round_train_loss)
+        self.train_round_metrics.append(round_train_metric)
+
+        metric_key = "train_accuracy" if self.local_model.task_type == "binary" else "train_rmse"
+
         return (
             self.local_model.get_parameters(),
-            len(self.trainloader),
-            {"epsilon": epsilon},
+            len(self.trainloader.dataset),   # better than len(dataloader)
+            {
+                "epsilon": float(epsilon),
+                "train_loss": float(round_train_loss),
+                metric_key: float(round_train_metric),
+            },
         )
 
+
     def evaluate(self, parameters, config):
-        """
-        Evaluates the local model on the validation data and returns the loss and accuracy.
+            print(f"[Client {self.cid}] evaluate, config: {config}")
+            self.local_model.set_parameters(parameters)
 
-        Args:
-            parameters: Parameters received from the server.
-            config: Configuration information.
+            # NEW: evaluate returns (loss, metrics_dict)
+            loss, metrics = self.local_model.evaluate(self.valloader, device=self.device)
 
-        Returns:
-            Tuple: Loss, number of validation examples, and accuracy information.
-        """
-        print(f"[Client {self.cid}] evaluate, config: {config}")
-        self.local_model.set_parameters(parameters)
-        loss, accuracy , auc = self.local_model.evaluate(
-            self.valloader, device=self.device
-        )
-        self.losses.append(loss)
-        self.accuracies.append(accuracy)
-        
-        return float(loss), len(self.valloader), {"accuracy": float(accuracy)}
+            self.losses.append(loss)
+
+            # Keep old behavior if accuracy exists
+            if isinstance(metrics, dict) and "accuracy" in metrics:
+                self.accuracies.append(metrics["accuracy"])
+
+            # Return metrics as-is (works for binary + regression)
+            # Flower expects: (loss, num_examples, metrics_dict)
+            return float(loss), len(self.valloader.dataset), {k: float(v) for k, v in (metrics or {}).items()}

medfl-2.0.6/MEDfl/LearningManager/model.py

@@ -0,0 +1,345 @@
+#!/usr/bin/env python3
+# froked from https://github.com/pythonlessons/mltu/blob/main/mltu/torch/model.py
+
+import typing
+from collections import OrderedDict
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from sklearn.metrics import accuracy_score, roc_auc_score, mean_squared_error, mean_absolute_error, r2_score
+
+
+from .utils import params
+
+
+class Model:
+    """
+    Model class for training and testing PyTorch neural networks.
+
+    Attributes:
+        model (torch.nn.Module): PyTorch neural network.
+        optimizer (torch.optim.Optimizer): PyTorch optimizer.
+        criterion (typing.Callable): Loss function.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimizer: torch.optim.Optimizer,
+        criterion: typing.Callable,
+        task_type: str = "binary"
+    ) -> None:
+        """
+        Initialize Model class with the specified model, optimizer, and criterion.
+
+        Args:
+            model (torch.nn.Module): PyTorch neural network.
+            optimizer (torch.optim.Optimizer): PyTorch optimizer.
+            criterion (typing.Callable): Loss function.
+        """
+        self.model = model
+        self.optimizer = optimizer
+        self.criterion = criterion
+        self.task_type = task_type
+        # Get device on which model is running
+        self.validate()
+
+    def validate(self) -> None:
+        """
+        Validate model and optimizer.
+        """
+        if not isinstance(self.model, torch.nn.Module):
+            raise TypeError("model argument must be a torch.nn.Module")
+
+        if not isinstance(self.optimizer, torch.optim.Optimizer):
+            raise TypeError(
+                "optimizer argument must be a torch.optim.Optimizer"
+            )
+
+    def get_parameters(self) -> List[np.ndarray]:
+        """
+        Get the parameters of the model as a list of NumPy arrays.
+
+        Returns:
+            List[np.ndarray]: The parameters of the model as a list of NumPy arrays.
+        """
+        return [val.detach().cpu().numpy() for val in self.model.state_dict().values()]
+
+    def set_parameters(self, parameters: List[np.ndarray]) -> None:
+        state_dict = self.model.state_dict()
+        keys = list(state_dict.keys())
+
+        device = next(self.model.parameters()).device  # model device
+        new_state_dict = OrderedDict()
+
+        for k, v in zip(keys, parameters):
+            ref = state_dict[k]  # reference tensor for dtype/shape
+            # v is numpy array -> tensor on correct device/dtype
+            new_state_dict[k] = torch.as_tensor(v, device=device, dtype=ref.dtype)
+
+        self.model.load_state_dict(new_state_dict, strict=True)
+
+    def train(self, train_loader, epoch, device, privacy_engine, diff_priv=False) -> float:
+        self.model.train()
+        epsilon = 0.0
+        losses = []
+        metrics = []
+
+        for i, (X_train, y_train) in enumerate(train_loader):
+            X_train = X_train.to(device)
+            y_train = y_train.to(device).view(-1)  # ensure (batch,)
+
+            # -------------------------
+            # HARD INPUT SANITY CHECKS
+            # -------------------------
+            if torch.isnan(X_train).any() or torch.isinf(X_train).any():
+                print(f"❌ Train batch {i}: X_train has NaN/Inf -> skipping batch")
+                continue
+            if torch.isnan(y_train).any() or torch.isinf(y_train).any():
+                print(f"❌ Train batch {i}: y_train has NaN/Inf -> skipping batch")
+                continue
+
+            # Optional: print scales once (helps diagnose exploding grads)
+            if i == 0 and epoch == 0:
+                try:
+                    print("Train scale check | X min/max:", X_train.min().item(), X_train.max().item(),
+                        "| y min/max:", y_train.min().item(), y_train.max().item())
+                except Exception:
+                    pass
+
+            self.optimizer.zero_grad(set_to_none=True)
+
+            # -------------------------
+            # FORWARD
+            # -------------------------
+            y_hat = self.model(X_train).view(-1)  # ensure (batch,)
+
+            if torch.isnan(y_hat).any() or torch.isinf(y_hat).any():
+                print(f"❌ Train batch {i}: y_hat has NaN/Inf -> skipping batch")
+                continue
+
+            # -------------------------
+            # LOSS
+            # -------------------------
+            loss = self.criterion(y_hat, y_train)
+
+            if torch.isnan(loss) or torch.isinf(loss):
+                print(f"❌ Train batch {i}: loss is NaN/Inf -> skipping batch")
+                continue
+
+            losses.append(float(loss.item()))
+
+            # -------------------------
+            # BACKWARD
+            # -------------------------
+            loss.backward()
+
+            # -------------------------
+            # GRADIENT SANITY CHECK
+            # If grads are NaN/Inf, DO NOT UPDATE WEIGHTS
+            # -------------------------
+            bad_grad = False
+            for name, p in self.model.named_parameters():
+                if p.grad is None:
+                    continue
+                if torch.isnan(p.grad).any() or torch.isinf(p.grad).any():
+                    print(
+                        f"❌ Train batch {i}: grad NaN/Inf in {name} "
+                        f"(NaN={torch.isnan(p.grad).sum().item()}, Inf={torch.isinf(p.grad).sum().item()}) "
+                        f"-> skipping update"
+                    )
+                    bad_grad = True
+                    break
+
+            if bad_grad:
+                self.optimizer.zero_grad(set_to_none=True)
+                continue
+
+            # -------------------------
+            # CLIP + STEP
+            # -------------------------
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            self.optimizer.step()
+
+            # -------------------------
+            # METRICS
+            # -------------------------
+            if self.task_type == "binary":
+                # Assumes y_hat is probability (Sigmoid model + BCELoss)
+                y_prob = y_hat.detach().cpu().numpy()
+                y_true = y_train.detach().cpu().numpy()
+                acc = accuracy_score(y_true, (y_prob >= 0.5).astype(int))
+                metrics.append(float(acc))
+
+            elif self.task_type == "regression":
+                rmse = torch.sqrt(torch.mean((y_hat - y_train) ** 2)).item()
+                metrics.append(float(rmse))
+
+            # -------------------------
+            # DP EPSILON
+            # -------------------------
+            if diff_priv:
+                epsilon = float(privacy_engine.get_epsilon(float(params["DELTA"])))
+
+            # -------------------------
+            # LOGGING
+            # -------------------------
+            if (i + 1) % 10 == 0 and len(losses) > 0:
+                if self.task_type == "binary" and len(metrics) > 0:
+                    print(
+                        f"\tTrain Epoch {epoch} | "
+                        f"Loss {np.mean(losses):.6f} | "
+                        f"Acc {np.mean(metrics) * 100:.2f}%"
+                        + (f" | ε={epsilon:.2f}" if diff_priv else "")
+                    )
+                elif self.task_type == "regression" and len(metrics) > 0:
+                    print(
+                        f"\tTrain Epoch {epoch} | "
+                        f"Loss {np.mean(losses):.6f} | "
+                        f"RMSE {np.mean(metrics):.6f}"
+                        + (f" | ε={epsilon:.2f}" if diff_priv else "")
+                    )
+
+        mean_train_loss = float(np.mean(losses)) if len(losses) else float("nan")
+        mean_train_metric = float(np.mean(metrics)) if len(metrics) else float("nan")
+
+        return epsilon, mean_train_loss, mean_train_metric
+
+
+
+    def evaluate(self, val_loader, device=torch.device("cpu")):
+        self.model.eval()
+        losses = []
+        y_true_all, y_pred_all = [], []
+
+        with torch.no_grad():
+            for i, (X_test, y_test) in enumerate(val_loader):
+                X_test = X_test.to(device)
+                y_test_t = y_test.to(device).view(-1)          # force 1D targets
+                y_hat_t = self.model(X_test).view(-1)          # force 1D preds
+
+                # Skip bad batches
+                if torch.isnan(y_hat_t).any() or torch.isinf(y_hat_t).any():
+                    print(f"❌ Eval batch {i}: y_hat NaN/Inf -> skipping")
+                    continue
+                if torch.isnan(y_test_t).any() or torch.isinf(y_test_t).any():
+                    print(f"❌ Eval batch {i}: y_test NaN/Inf -> skipping")
+                    continue
+
+                crit = self.criterion.to(y_hat_t.device) if hasattr(self.criterion, "to") else self.criterion
+                loss_t = crit(y_hat_t, y_test_t)
+
+                if torch.isnan(loss_t) or torch.isinf(loss_t):
+                    print(f"❌ Eval batch {i}: loss NaN/Inf -> skipping")
+                    continue
+
+                losses.append(float(loss_t.item()))
+                # IMPORTANT: append 1D numpy arrays (never scalars)
+                y_true_all.append(np.atleast_1d(y_test_t.detach().cpu().numpy()))
+                y_pred_all.append(np.atleast_1d(y_hat_t.detach().cpu().numpy()))
+
+        if len(losses) == 0 or len(y_true_all) == 0:
+            return float("nan"), {}
+
+        y_true = np.concatenate(y_true_all, axis=0)
+        y_pred = np.concatenate(y_pred_all, axis=0)
+
+        mean_loss = float(np.mean(losses))
+
+        if self.task_type == "binary":
+            acc = accuracy_score(y_true, (y_pred >= 0.5).astype(int))
+            auc = roc_auc_score(y_true, y_pred) if len(np.unique(y_true)) > 1 else float("nan")
+            return mean_loss, {"accuracy": float(acc), "auc": float(auc)}
+
+        if self.task_type == "regression":
+            mask = np.isfinite(y_true) & np.isfinite(y_pred)
+            if mask.sum() == 0:
+                return mean_loss, {
+                    "rmse": float("nan"),
+                    "mae": float("nan"),
+                    "r2": float("nan"),
+                    "adj_r2": float("nan"),
+                }
+
+            y_true_c = y_true[mask]
+            y_pred_c = y_pred[mask]
+
+            mse = mean_squared_error(y_true_c, y_pred_c)
+            mae = mean_absolute_error(y_true_c, y_pred_c)
+            r2 = r2_score(y_true_c, y_pred_c)
+            rmse = float(np.sqrt(mse))
+
+            # ---- Adjusted R² ----
+            n = int(len(y_true_c))
+
+            # Determine p (number of predictors)
+            p = None
+            try:
+                # Works for Linear Regressor and most MLPs (first Linear layer)
+                for m in self.model.modules():
+                    if isinstance(m, torch.nn.Linear):
+                        p = int(m.in_features)
+                        break
+            except Exception:
+                p = None
+
+            if p is None or n <= p + 1:
+                adj_r2 = float("nan")
+            else:
+                adj_r2 = 1.0 - (1.0 - float(r2)) * (n - 1) / (n - p - 1)
+            
+            print("adj_r2 check | n =", n, "| p =", p, "| condition n>p+1 =", n > p + 1)
+
+
+            return mean_loss, {
+                "rmse": float(rmse),
+                "mae": float(mae),
+                "r2": float(r2),
+                "adj_r2": float(adj_r2),
+            }
+
+        return mean_loss, {}
+
+
+    @staticmethod
+    def save_model(model , model_name:str):
+        """
+        Saves a PyTorch model to a file.
+
+        Args:
+            model (torch.nn.Module): PyTorch model to be saved.
+            model_name (str): Name of the model file.
+
+        Raises:
+            Exception: If there is an issue during the saving process.
+
+        Returns:
+            None
+        """
+        try:
+            torch.save(model, '../../notebooks/.ipynb_checkpoints/trainedModels/' + model_name + ".pth")
+        except Exception as e:
+            raise Exception(f"Error saving the model: {str(e)}")
+    
+    @staticmethod
+    def load_model(model_path: str):
+        """
+        Loads a PyTorch model from a file.
+
+        Args:
+            model_path (str): Path to the model file to be loaded.
+
+        Returns:
+            torch.nn.Module: Loaded PyTorch model.
+        """
+        # Ensure models are loaded onto the CPU when CUDA is not available
+        torch_kwargs = {"weights_only": False}
+        if torch.cuda.is_available():
+            loaded_model = torch.load(model_path , **torch_kwargs)
+        else:
+            loaded_model = torch.load(model_path, map_location=torch.device('cpu') , **torch_kwargs)
+        return loaded_model
+
+

medfl-2.0.6/MEDfl/LearningManager/server.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+
+import copy
+from typing import Dict, Optional, Tuple
+
+import flwr as fl
+import torch
+
+from torch.utils.data import ConcatDataset, DataLoader  # ✅ NEW
+
+from .client import FlowerClient
+from .federated_dataset import FederatedDataset
+from .model import Model
+from .strategy import Strategy
+
+
+class FlowerServer:
+    """
+    A class representing the central server for Federated Learning using Flower.
+    """
+
+    def __init__(
+        self,
+        global_model: Model,
+        strategy: Strategy,
+        num_rounds: int,
+        num_clients: int,
+        fed_dataset: FederatedDataset,
+        diff_privacy: bool = False,
+        client_resources: Optional[Dict[str, float]] = {'num_cpus': 1, 'num_gpus': 0.0}
+    ) -> None:
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.global_model = global_model
+        self.params = global_model.get_parameters()
+        self.global_model.model = global_model.model.to(self.device)
+        self.num_rounds = num_rounds
+        self.num_clients = num_clients
+        self.fed_dataset = fed_dataset
+        self.strategy = strategy
+        self.client_resources = client_resources
+
+        setattr(self.strategy.strategy_object, "min_available_clients", self.num_clients)
+        setattr(self.strategy.strategy_object, "initial_parameters", fl.common.ndarrays_to_parameters(self.params))
+        setattr(self.strategy.strategy_object, "evaluate_fn", self.evaluate)
+
+        self.diff_priv = diff_privacy
+
+        # ✅ tracking
+        self.accuracies = []
+        self.losses = []
+        self.auc = []
+        self.rmse = []
+        self.mae = []
+        self.r2 = []
+        self.adj_r2 = []
+
+        self.train_losses = []
+        self.train_accuracies = []
+        self.train_auc = []
+        self.train_rmse = []
+        self.train_mae = []
+        self.train_r2 = []
+        self.train_adj_r2 = []
+
+        self.flower_clients = []
+        self.device = torch.device("cpu")
+        self.global_model = global_model
+
+        # move model to CPU
+        self.global_model.model = self.global_model.model.to(self.device)
+
+        # also move optimizer state to CPU (important if optimizer had any tensors)
+        for state in self.global_model.optimizer.state.values():
+            for k, v in state.items():
+                if torch.is_tensor(v):
+                    state[k] = v.detach().cpu()
+
+        # if criterion has tensor buffers (e.g., pos_weight), ensure CPU too
+        if hasattr(self.global_model.criterion, "pos_weight") and self.global_model.criterion.pos_weight is not None:
+            self.global_model.criterion.pos_weight = self.global_model.criterion.pos_weight.detach().cpu()
+        self.validate()
+
+    def validate(self) -> None:
+        if not isinstance(self.global_model, Model):
+            raise TypeError("global_model argument must be a Model instance")
+        if not isinstance(self.num_clients, int):
+            raise TypeError("num_clients argument must be an int")
+        if not isinstance(self.num_rounds, int):
+            raise TypeError("num_rounds argument must be an int")
+        if not isinstance(self.diff_priv, bool):
+            raise TypeError("diff_priv argument must be a bool")
+
+    def client_fn(self, cid) -> FlowerClient:
+        # NOTE: you might want to replace %4 with device_count() later
+        device = torch.device(f"cuda:{int(cid) % 4}" if torch.cuda.is_available() else "cpu")
+
+        client_model = copy.deepcopy(self.global_model)
+
+        trainloader = self.fed_dataset.trainloaders[int(cid)]
+        valloader = self.fed_dataset.valloaders[int(cid)]
+
+        client_model = copy.deepcopy(self.global_model)
+        client = FlowerClient(cid, client_model, trainloader, valloader, self.diff_priv)
+        self.flower_clients.append(client)
+        return client
+
+    # ✅ NEW helper: build one global loader from all client loaders (simulation)
+    def _concat_loader(self, loaders, shuffle: bool = False) -> DataLoader:
+        # Use each loader's dataset and build one union dataset
+        ds = ConcatDataset([ldr.dataset for ldr in loaders])
+
+        # Try to reuse an existing batch_size if possible; fallback to 64
+        batch_size = getattr(loaders[0], "batch_size", None) or 64
+        num_workers = getattr(loaders[0], "num_workers", 0)
+        pin_memory = getattr(loaders[0], "pin_memory", False)
+
+        return DataLoader(
+            ds,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            num_workers=num_workers,
+            pin_memory=pin_memory,
+        )
+
+    def evaluate(
+        self,
+        server_round: int,
+        parameters: fl.common.NDArrays,
+        config: Dict[str, fl.common.Scalar],
+    ) -> Optional[Tuple[float, Dict[str, fl.common.Scalar]]]:
+
+        # Update model with latest weights
+        self.global_model.set_parameters(parameters)
+        self.last_parameters = parameters
+
+
+        # ✅ GLOBAL loaders across ALL clients (simulation)
+        global_valloader = self._concat_loader(self.fed_dataset.valloaders, shuffle=False)
+        global_trainloader = self._concat_loader(self.fed_dataset.trainloaders, shuffle=False)
+
+        # ---- Validation (global) ----
+        val_loss, val_metrics = self.global_model.evaluate(global_valloader, self.device)
+        self.losses.append(val_loss)
+
+        if isinstance(val_metrics, dict):
+            if "accuracy" in val_metrics: self.accuracies.append(val_metrics["accuracy"])
+            if "auc" in val_metrics: self.auc.append(val_metrics["auc"])
+            if "rmse" in val_metrics: self.rmse.append(val_metrics["rmse"])
+            if "mae" in val_metrics: self.mae.append(val_metrics["mae"])
+            if "r2" in val_metrics: self.r2.append(val_metrics["r2"])
+            if "adj_r2" in val_metrics: self.adj_r2.append(val_metrics["adj_r2"])
+
+        # ---- Train (global) ----
+        train_loss, train_metrics = self.global_model.evaluate(global_trainloader, self.device)
+        self.train_losses.append(train_loss)
+
+        if isinstance(train_metrics, dict):
+            if "accuracy" in train_metrics: self.train_accuracies.append(train_metrics["accuracy"])
+            if "auc" in train_metrics: self.train_auc.append(train_metrics["auc"])
+            if "rmse" in train_metrics: self.train_rmse.append(train_metrics["rmse"])
+            if "mae" in train_metrics: self.train_mae.append(train_metrics["mae"])
+            if "r2" in train_metrics: self.train_r2.append(train_metrics["r2"])
+            if "adj_r2" in train_metrics: self.train_adj_r2.append(train_metrics["adj_r2"])
+
+        # Return validation loss + a metrics dict that also includes train_*
+        out_metrics = dict(val_metrics or {})
+        out_metrics["train_loss"] = float(train_loss) if train_loss is not None else float("nan")
+
+        if isinstance(train_metrics, dict):
+            for k, v in train_metrics.items():
+                out_metrics[f"train_{k}"] = float(v) if v is not None else float("nan")
+
+        return val_loss, out_metrics
+
+    def score_each_client(self, parameters=None, split: str = "val"):
+        """
+        Evaluate the (final) global model on each client's data.
+        split: "val" or "train"
+        """
+        if parameters is not None:
+            self.global_model.set_parameters(parameters)
+
+        loaders = self.fed_dataset.valloaders if split == "val" else self.fed_dataset.trainloaders
+
+        results = {}
+        for cid, loader in enumerate(loaders):
+            loss, metrics = self.global_model.evaluate(loader, self.device)
+            results[str(cid)] = {
+                "num_examples": len(loader.dataset),
+                "loss": float(loss),
+                **{k: float(v) for k, v in (metrics or {}).items()},
+            }
+        return results
+
+
+
+    def run(self) -> None:
+        ray_init_args = {"include_dashboard": False, "object_store_memory": 78643200 , 
+                        }
+        self.fed_dataset.eng = None
+
+        history = fl.simulation.start_simulation(
+            client_fn=self.client_fn,
+            num_clients=self.num_clients,
+            config=fl.server.ServerConfig(self.num_rounds),
+            strategy=self.strategy.strategy_object,
+            ray_init_args=ray_init_args,
+            client_resources=self.client_resources
+        )
+        return history