MEDfl 2.0.9__tar.gz → 2.0.10.dev0__tar.gz

medfl-2.0.10.dev0/MEDfl/LearningManager/client.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+import logging
+import sys
+
+import flwr as fl
+from opacus import PrivacyEngine
+from torch.utils.data import DataLoader
+
+from .model import Model
+from .utils import params
+import torch
+import numpy as np
+
+
+# ---------------------------------------------------------------------------
+# Per-worker logger setup
+# ---------------------------------------------------------------------------
+# Ray workers are separate processes. We configure this logger inside the
+# worker so that messages go to stdout — Ray then forwards stdout to the
+# driver when log_to_driver=True is set in ray_init_args.
+# ---------------------------------------------------------------------------
+def _get_logger() -> logging.Logger:
+    logger = logging.getLogger("MEDfl.client")
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(
+            logging.Formatter(
+                "[%(asctime)s][%(name)s][%(levelname)s] %(message)s",
+                datefmt="%H:%M:%S",
+            )
+        )
+        logger.addHandler(handler)
+        logger.setLevel(logging.DEBUG)
+        logger.propagate = False  # avoid double-printing on driver
+    return logger
+
+
+class FlowerClient(fl.client.NumPyClient):
+    """
+    FlowerClient class for creating MEDfl clients.
+    """
+
+    def __init__(
+        self,
+        cid: str,
+        local_model: Model,
+        trainloader: DataLoader,
+        valloader: DataLoader,
+        diff_priv: bool = params["diff_privacy"],
+    ):
+        self.cid = cid
+        self.local_model = local_model
+        self.trainloader = trainloader
+        self.valloader = valloader
+        self.log = _get_logger()
+
+        self.device = torch.device("cpu")
+        self.local_model.model.to(self.device)
+
+        self.privacy_engine = PrivacyEngine(secure_mode=False)
+        self.diff_priv = diff_priv
+        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(),
+                optimizer=self.local_model.optimizer,
+                data_loader=self.trainloader,
+                epochs=params["train_epochs"],
+                target_epsilon=float(params["EPSILON"]),
+                target_delta=float(params["DELTA"]),
+                max_grad_norm=params["MAX_GRAD_NORM"],
+            )
+            setattr(self.local_model, "model", model)
+            setattr(self.local_model, "optimizer", optimizer)
+        self.validate()
+
+    # ------------------------------------------------------------------
+    def validate(self):
+        """Validates cid, local_model, trainloader, valloader."""
+        if not isinstance(self.cid, str):
+            raise TypeError("cid argument must be a string")
+        if not isinstance(self.local_model, Model):
+            raise TypeError("local_model argument must be a MEDfl.LearningManager.model.Model")
+        if not isinstance(self.trainloader, DataLoader):
+            raise TypeError("trainloader argument must be a torch.utils.data.dataloader")
+        if not isinstance(self.valloader, DataLoader):
+            raise TypeError("valloader argument must be a torch.utils.data.dataloader")
+        if not isinstance(self.diff_priv, bool):
+            raise TypeError("diff_priv argument must be a bool")
+
+    # ------------------------------------------------------------------
+    def get_parameters(self, config):
+        self.log.info(f"[Client {self.cid}] get_parameters")
+        return self.local_model.get_parameters()
+
+    # ------------------------------------------------------------------
+    def fit(self, parameters, config):
+        self.log.info(f"[Client {self.cid}] fit | config: {config}")
+
+        self.local_model.set_parameters(parameters)
+
+        epoch_losses, epoch_metrics = [], []
+
+        for ep in range(params["train_epochs"]):
+            epsilon, tr_loss, tr_metric = self.local_model.train(
+                self.trainloader,
+                epoch=ep,
+                device=self.device,
+                privacy_engine=self.privacy_engine,
+                diff_priv=self.diff_priv,
+            )
+            self.epsilons.append(epsilon)
+            epoch_losses.append(tr_loss)
+            epoch_metrics.append(tr_metric)
+        
+        avg_grad_norm = float(np.nanmean(self.local_model.grad_norms)) \
+                        if hasattr(self.local_model, "grad_norms") and self.local_model.grad_norms \
+                        else float("nan")
+        self.local_model.grad_norms = []  # reset for next round
+
+
+        self.log.info(
+            f"[Client {self.cid}] fit done | "
+            f"epsilon={epsilon:.4f} | "
+            f"avg_loss={float(np.nanmean(epoch_losses)):.4f}"
+        )
+
+        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.dataset),
+            {
+                "epsilon": float(epsilon),
+                "train_loss": float(round_train_loss),
+                metric_key: float(round_train_metric),
+                "grad_norm": avg_grad_norm,   # <-- NEW
+            },
+        )
+
+    # ------------------------------------------------------------------
+    def evaluate(self, parameters, config):
+        self.log.info(f"[Client {self.cid}] evaluate | config: {config}")
+        self.local_model.set_parameters(parameters)
+
+        loss, metrics = self.local_model.evaluate(self.valloader, device=self.device)
+
+        self.losses.append(loss)
+
+        if isinstance(metrics, dict) and "accuracy" in metrics:
+            self.accuracies.append(metrics["accuracy"])
+
+        self.log.info(
+            f"[Client {self.cid}] evaluate done | "
+            f"loss={loss:.4f} | metrics={metrics}"
+        )
+
+        return float(loss), len(self.valloader.dataset), {k: float(v) for k, v in (metrics or {}).items()}

{medfl-2.0.9 → medfl-2.0.10.dev0}/MEDfl/LearningManager/model.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
-# froked from https://github.com/pythonlessons/mltu/blob/main/mltu/torch/model.py
 
+import logging
+import sys
 import typing
 from collections import OrderedDict
 from typing import Dict, List, Optional, Tuple
@@ -8,20 +9,39 @@ 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 sklearn.metrics import (
+    accuracy_score, roc_auc_score,
+    mean_squared_error, mean_absolute_error, r2_score,
+)
 
 from .utils import params
 
 
+# ---------------------------------------------------------------------------
+# Per-worker logger — same pattern as client.py so Ray forwards to notebook
+# ---------------------------------------------------------------------------
+def _get_logger() -> logging.Logger:
+    logger = logging.getLogger("MEDfl.model")
+    if not logger.handlers:
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(
+            logging.Formatter(
+                "[%(asctime)s][%(name)s][%(levelname)s] %(message)s",
+                datefmt="%H:%M:%S",
+            )
+        )
+        logger.addHandler(handler)
+        logger.setLevel(logging.DEBUG)
+        logger.propagate = False
+    return logger
+
+
+_log = _get_logger()
+
+
 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__(
@@ -29,56 +49,31 @@ class Model:
         model: torch.nn.Module,
         optimizer: torch.optim.Optimizer,
         criterion: typing.Callable,
-        task_type: str = "binary"
+        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"
-            )
+            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
+        device = next(self.model.parameters()).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
+            ref = state_dict[k]
             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:
@@ -89,66 +84,51 @@ class Model:
 
         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,)
+            y_train = y_train.to(device).view(-1)
 
-            # -------------------------
-            # 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")
+                _log.warning(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")
+                _log.warning(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())
+                    _log.debug(
+                        f"Train scale check | "
+                        f"X min/max: {X_train.min().item():.4f} / {X_train.max().item():.4f} | "
+                        f"y min/max: {y_train.min().item():.4f} / {y_train.max().item():.4f}"
+                    )
                 except Exception:
                     pass
 
             self.optimizer.zero_grad(set_to_none=True)
 
-            # -------------------------
-            # FORWARD
-            # -------------------------
-            y_hat = self.model(X_train).view(-1)  # ensure (batch,)
+            y_hat = self.model(X_train).view(-1)
 
             if torch.isnan(y_hat).any() or torch.isinf(y_hat).any():
-                print(f"❌ Train batch {i}: y_hat has NaN/Inf -> skipping batch")
+                _log.warning(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")
+                _log.warning(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"
+                    _log.warning(
+                        f"Train batch {i}: grad NaN/Inf in {name} "
+                        f"(NaN={torch.isnan(p.grad).sum().item()}, "
+                        f"Inf={torch.isinf(p.grad).sum().item()}) -> skipping update"
                     )
                     bad_grad = True
                     break
@@ -157,46 +137,32 @@ class Model:
                 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} | "
+                    _log.info(
+                        f"Train Epoch {epoch} | batch {i+1} | "
                         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} | "
+                    _log.info(
+                        f"Train Epoch {epoch} | batch {i+1} | "
                         f"Loss {np.mean(losses):.6f} | "
                         f"RMSE {np.mean(metrics):.6f}"
                         + (f" | ε={epsilon:.2f}" if diff_priv else "")
@@ -205,9 +171,17 @@ class Model:
         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
-
+        total_grad_norm = 0.0
+        for p in self.model.parameters():
+            if p.grad is not None:
+                total_grad_norm += p.grad.detach().norm(2).item() ** 2
+        total_grad_norm = total_grad_norm ** 0.5
+        # store it
+        if not hasattr(self, "grad_norms"):
+            self.grad_norms = []
+        self.grad_norms.append(total_grad_norm)
 
+        return epsilon, mean_train_loss, mean_train_metric
 
     def evaluate(self, val_loader, device=torch.device("cpu")):
         self.model.eval()
@@ -217,26 +191,24 @@ class Model:
         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
+                y_test_t = y_test.to(device).view(-1)
+                y_hat_t = self.model(X_test).view(-1)
 
-                # 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")
+                    _log.warning(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")
+                    _log.warning(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")
+                    _log.warning(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()))
 
@@ -245,7 +217,6 @@ class Model:
 
         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":
@@ -256,28 +227,19 @@ class Model:
         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"),
-                }
+                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)
@@ -289,9 +251,8 @@ class Model:
                 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)
 
+            _log.debug(f"adj_r2 check | n={n} | p={p} | condition n>p+1={n > p + 1}")
 
             return mean_loss, {
                 "rmse": float(rmse),
@@ -302,44 +263,18 @@ class Model:
 
         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
-        """
+    def save_model(model, model_name: str):
         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)
+            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
-
-
+            loaded_model = torch.load(model_path, map_location=torch.device('cpu'), **torch_kwargs)
+        return loaded_model

{medfl-2.0.9 → medfl-2.0.10.dev0}/MEDfl/LearningManager/params.yaml

@@ -1,8 +1,8 @@
 DELTA: 1.0e-05
 EPSILON: 5.0
 MAX_GRAD_NORM: 1.0
-diff_privacy: true
-lr: 0.001
+diff_privacy: false
+lr: 0.01
 min_evalclient: 2
 num_rounds: 20
 optimizer: SGD
@@ -10,5 +10,5 @@ path_to_master_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/
 path_to_test_csv: /home/local/USHERBROOKE/saho6810/MEDfl/code/MEDfl/notebooks/data/masterDataSet/Mimic_train.csv
 task: BinaryClassification
 test_batch_size: 1
-train_batch_size: 32
-train_epochs: 20
+train_batch_size: 6
+train_epochs: 1