janky but working

Author: nerdai

fl4health/clients/nnunet_client.py

@@ -1,4 +1,5 @@
 import gc
+import copy
 import logging
 import os
 import pickle
@@ -186,6 +187,7 @@ def __init__(
         self.nnunet_trainer_class = nnunet_trainer_class
         self.nnunet_trainer_class_kwargs = nnunet_trainer_class_kwargs
         self.nnunet_trainer: nnUNetTrainer
+
         self.nnunet_config: NnunetConfig
         self.plans: dict[str, Any] | None = None
         self.steps_per_round: int  # N steps per server round
@@ -227,7 +229,7 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
 
         # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler
         # Clear gradients from optimizer if they exist
-        self.optimizers["global"].zero_grad()
+        self.optimizers["local"].zero_grad()
 
         # Call user defined methods to get predictions and compute loss
         preds, features = self.predict(input)
@@ -239,11 +241,11 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
         scaled_backward_loss.backward()
 
         # Rescale gradients then clip based on specified norm
-        self.grad_scaler.unscale_(self.optimizers["global"])
+        self.grad_scaler.unscale_(self.optimizers["local"])
         self.transform_gradients(losses)
 
         # Update parameters and scaler
-        self.grad_scaler.step(self.optimizers["global"])
+        self.grad_scaler.step(self.optimizers["local"])
         self.grad_scaler.update()
 
         return losses, preds
@@ -314,7 +316,11 @@ def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
         return train_loader, val_loader
 
     def get_model(self, config: Config) -> nn.Module:
-        return self.nnunet_trainer.network
+        for_global = config.get("for_global", False)
+        if for_global:
+            return copy.deepcopy(self.nnunet_trainer.network)
+        else:
+            return self.nnunet_trainer.network
 
     def get_criterion(self, config: Config) -> _Loss:
         if isinstance(self.nnunet_trainer.loss, DeepSupervisionWrapper):
@@ -608,7 +614,6 @@ def setup_client(self, config: Config) -> None:
     def _special_predict(
         self, model: torch.nn.Module, input: torch.Tensor
     ) -> tuple[TorchPredType, dict[str, torch.Tensor]]:
-        model.train()
         if isinstance(input, torch.Tensor):
             # If device type is cuda, nnUNet defaults to mixed precision forward pass
             if self.device.type == "cuda":
@@ -770,8 +775,14 @@ def update_metric_manager(
             target (TorchTargetType): the targets generated by the dataloader to evaluate the preds with
             metric_manager (MetricManager): the metric manager to update
         """
+        preds = {k: v for k, v in preds.items() if "local" in k}
+        # remove prefix
+        preds = {k.replace(f"local-", ""): v for k, v in preds.items()}
+
         if len(preds) > 1:
             # for nnunet the first pred in the output list is the main one
+            log(DEBUG, f"preds keys: {preds.keys()}")
+
             m_pred = convert_deep_supervision_dict_to_list(preds)[0]
 
         if isinstance(target, torch.Tensor):
@@ -828,7 +839,7 @@ def get_client_specific_logs(
         logging_mode: LoggingMode,
     ) -> tuple[str, list[tuple[LogLevel, str]]]:
         if logging_mode == LoggingMode.TRAIN:
-            lr = float(self.optimizers["global"].param_groups[0]["lr"])
+            lr = float(self.optimizers["local"].param_groups[0]["lr"])
             if current_epoch is None:
                 # Assume training by steps
                 return f"Initial LR {lr}", []
@@ -838,7 +849,7 @@ def get_client_specific_logs(
             return "", []
 
     def get_client_specific_reports(self) -> dict[str, Any]:
-        return {"learning_rate": float(self.optimizers["global"].param_groups[0]["lr"])}
+        return {"learning_rate": float(self.optimizers["local"].param_groups[0]["lr"])}
 
     @use_default_signal_handlers  # Experiment planner spawns a process I think
     def get_properties(self, config: Config) -> dict[str, Scalar]:
@@ -942,12 +953,13 @@ def update_before_train(self, current_server_round: int) -> None:
             # freeze before the first pass, gc.collect has to check all those variables
             gc.freeze()
 
-    def transform_gradients(self, losses: TrainingLosses) -> None:
+    def transform_gradients(self, losses: TrainingLosses, model: nn.Module | None = None) -> None:
         """
         Apply the gradient clipping performed by the default nnunet trainer. This is the default behavior for
         nnunet 2.5.1
 
         Args:
             losses (TrainingLosses): Not used for this transformation.
         """
-        nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+        model = model if model else self.model
+        nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm)

fl4health/mixins/personalized/ditto.py

@@ -93,6 +93,7 @@ def get_global_model(self, config: Config) -> nn.Module:
         Returns:
             nn.Module: The PyTorch model serving as the global model for Ditto
         """
+        config["for_global"] = True
         return self.get_model(config).to(self.device)
 
     def set_optimizer(self, config: Config) -> None:
@@ -315,7 +316,24 @@ def predict(
         # TODO: Perhaps loosen this at a later date.
         # assert isinstance(global_preds, torch.Tensor)
         # assert isinstance(local_preds, torch.Tensor)
-        return {"global": global_preds, "local": local_preds}, {}
+        if isinstance(global_preds, torch.Tensor) and isinstance(local_preds, torch.Tensor):
+            return {"global": global_preds, "local": local_preds}, {}
+        elif isinstance(global_preds, dict) and isinstance(local_preds, dict):
+            retval = {f"global-{k}": v for k, v in global_preds.items()}
+            retval.update(**{f"local-{k}": v for k, v in local_preds.items()})
+            return retval, {}
+        else:
+            raise ValueError(f"Unsupported pred type: {type(global_preds)}.")
+
+    def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]):
+        if kind not in ["global", "local"]:
+            raise ValueError("Unsupported kind of prediction. Must be 'global' or 'local'.")
+
+        # filter
+        retval = {k: v for k, v in preds.items() if kind in k}
+        # remove prefix
+        retval = {k.replace(f"{kind}-", ""): v for k, v in retval.items()}
+        return retval
 
     def compute_loss_and_additional_losses(
         self,
@@ -338,20 +356,23 @@ def compute_loss_and_additional_losses(
             - A dictionary with ``local_loss``, ``global_loss`` as additionally reported loss values.
         """
 
+        global_preds = self._extract_pred(kind="global", preds=preds)
+        local_preds = self._extract_pred(kind="local", preds=preds)
+
         # Compute global model vanilla loss
 
         if hasattr(self, "_special_compute_loss_and_additional_losses"):
             log(INFO, "Using '_special_compute_loss_and_additional_losses' to compute loss")
-            global_loss, _ = self._special_compute_loss_and_additional_losses(preds["global"], features, target)
+            global_loss, _ = self._special_compute_loss_and_additional_losses(global_preds, features, target)
 
             # Compute local model loss + ditto constraint term
-            local_loss, _ = self._special_compute_loss_and_additional_losses(preds["local"], features, target)
+            local_loss, _ = self._special_compute_loss_and_additional_losses(local_preds, features, target)
 
         else:
-            global_loss = self.criterion(preds["global"], target)
+            global_loss = self.criterion(global_preds, target)
 
             # Compute local model loss + ditto constraint term
-            local_loss = self.criterion(preds["local"], target)
+            local_loss = self.criterion(local_preds, target)
 
         additional_losses = {"local_loss": local_loss.clone(), "global_loss": global_loss}