From 55efba1a37c9519b19c920bd860032a0d9c018de Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 21 May 2025 11:49:05 -0400 Subject: [PATCH 01/32] use 'local' as key for optimizer --- fl4health/clients/nnunet_client.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index ecf9b57bf..2aa850612 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -228,7 +228,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) @@ -240,11 +240,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 @@ -800,7 +800,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}", [] @@ -810,7 +810,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]: From 73e302ddf227cc7bd2126aa44a277c3687e2d7dc Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 21 May 2025 12:08:10 -0400 Subject: [PATCH 02/32] add FOR_GLOBAL_MODEL_KEY in config.utils --- fl4health/clients/nnunet_client.py | 10 +++++++++- fl4health/mixins/personalized/ditto.py | 3 ++- fl4health/utils/config.py | 2 ++ 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 2aa850612..5ba605976 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -1,3 +1,4 @@ +import copy import gc import logging import os @@ -27,7 +28,7 @@ from fl4health.metrics.base_metrics import Metric from fl4health.metrics.metric_managers import MetricManager from fl4health.reporting.base_reporter import BaseReporter -from fl4health.utils.config import narrow_dict_type +from fl4health.utils.config import FOR_GLOBAL_MODEL_KEY, narrow_dict_type from fl4health.utils.logging import LoggingMode from fl4health.utils.losses import LossMeterType, TrainingLosses from fl4health.utils.nnunet_utils import ( @@ -315,6 +316,13 @@ def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]: return train_loader, val_loader def get_model(self, config: Config) -> nn.Module: + + for_global = config.get(FOR_GLOBAL_MODEL_KEY, False) + if for_global: + return copy.deepcopy(self.nnunet_trainer.network) + else: + return self.nnunet_trainer.network + return self.nnunet_trainer.network def get_criterion(self, config: Config) -> _Loss: diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index 7ceae19c7..a2d9df064 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -15,7 +15,7 @@ from fl4health.mixins.core_protocols import BasicClientProtocolPreSetup from fl4health.mixins.personalized.utils import ensure_protocol_compliance from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger -from fl4health.utils.config import narrow_dict_type +from fl4health.utils.config import FOR_GLOBAL_MODEL_KEY, narrow_dict_type from fl4health.utils.losses import EvaluationLosses, TrainingLosses from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType @@ -175,6 +175,7 @@ def get_global_model(self: DittoPersonalizedProtocol, config: Config) -> nn.Modu Returns: nn.Module: The PyTorch model serving as the global model for Ditto """ + config[FOR_GLOBAL_MODEL_KEY] = True return self.get_model(config).to(self.device) @ensure_protocol_compliance diff --git a/fl4health/utils/config.py b/fl4health/utils/config.py index 7a8becaf9..8beb2b70d 100644 --- a/fl4health/utils/config.py +++ b/fl4health/utils/config.py @@ -8,6 +8,8 @@ "batch_size": int, } +FOR_GLOBAL_MODEL_KEY = "__for_global_model" + T = TypeVar("T") From 0641991f731458460432dbdae39dc55ea0c2f761 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 21 May 2025 12:09:40 -0400 Subject: [PATCH 03/32] add FOR_GLOBAL_MODEL_KEY in config.utils --- fl4health/clients/nnunet_client.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 5ba605976..853aa5ec0 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -323,8 +323,6 @@ def get_model(self, config: Config) -> nn.Module: else: return self.nnunet_trainer.network - return self.nnunet_trainer.network - def get_criterion(self, config: Config) -> _Loss: if isinstance(self.nnunet_trainer.loss, DeepSupervisionWrapper): self.reports_manager.report({"Criterion": self.nnunet_trainer.loss.loss.__class__.__name__}) From 6ef3b2fd8f7e8857ba15526a331bc0e38d4fe9a0 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 21 May 2025 12:59:23 -0400 Subject: [PATCH 04/32] wip --- fl4health/clients/nnunet_client.py | 76 ++++++++++++++++++------------ 1 file changed, 46 insertions(+), 30 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 853aa5ec0..9d92444f9 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -612,27 +612,14 @@ def setup_client(self, config: Config) -> None: # We have to call parent method after setting up nnunet trainer super().setup_client(config) - def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]: - """ - Generate model outputs. Overridden because nnunets outputs lists when deep supervision is on so we have to - reformat the output into dicts. - - Additionally if device type is cuda, loss computed in mixed precision. - - Args: - input (TorchInputType): The model inputs - - Returns: - tuple[TorchPredType, dict[str, torch.Tensor]]: A tuple in which the first element model outputs indexed by - name. The second element is unused by this subclass and therefore is always an empty dict - """ + def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]: if isinstance(input, torch.Tensor): # If device type is cuda, nnUNet defaults to mixed precision forward pass if self.device.type == "cuda": with torch.autocast(self.device.type, enabled=True): - output = self.model(input) + output = model(input) else: - output = self.model(input) + output = model(input) else: raise TypeError('"input" must be of type torch.Tensor for nnUNetClient') @@ -648,26 +635,28 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch "Was expecting nnunet model output to be either a torch.Tensor or a list/tuple of torch.Tensors" ) - def compute_loss_and_additional_losses( - self, - preds: TorchPredType, - features: dict[str, torch.Tensor], - target: TorchTargetType, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]: """ - Checks the pred and target types and computes the loss. If device type is cuda, loss computed in mixed - precision. + Generate model outputs. Overridden because nnunets outputs lists when deep supervision is on so we have to + reformat the output into dicts. + + Additionally if device type is cuda, loss computed in mixed precision. Args: - preds (TorchPredType): Dictionary of model output tensors indexed by name - features (dict[str, torch.Tensor]): Not used by this subclass - target (TorchTargetType): The targets to evaluate the predictions with. If multiple prediction tensors - are given, target must be a dictionary with the same number of tensors + input (TorchInputType): The model inputs Returns: - tuple[torch.Tensor, dict[str, torch.Tensor] | None]: A tuple where the first element is the loss and the - second element is an optional additional loss + tuple[TorchPredType, dict[str, torch.Tensor]]: A tuple in which the first element model outputs indexed by + name. The second element is unused by this subclass and therefore is always an empty dict """ + return self._predict(self.model, input) + + def _special_compute_loss_and_additional_losses( + self, + preds: TorchPredType, + features: dict[str, torch.Tensor], + target: TorchTargetType, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: # If deep supervision is turned on we must convert loss and target dicts into lists loss_preds = prepare_loss_arg(preds) loss_targets = prepare_loss_arg(target) @@ -692,6 +681,28 @@ def compute_loss_and_additional_losses( return loss + def compute_loss_and_additional_losses( + self, + preds: TorchPredType, + features: dict[str, torch.Tensor], + target: TorchTargetType, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + """ + Checks the pred and target types and computes the loss. If device type is cuda, loss computed in mixed + precision. + + Args: + preds (TorchPredType): Dictionary of model output tensors indexed by name + features (dict[str, torch.Tensor]): Not used by this subclass + target (TorchTargetType): The targets to evaluate the predictions with. If multiple prediction tensors + are given, target must be a dictionary with the same number of tensors + + Returns: + tuple[torch.Tensor, dict[str, torch.Tensor] | None]: A tuple where the first element is the loss and the + second element is an optional additional loss + """ + return self._special_compute_loss_and_additional_losses(preds, features, target) + def mask_data(self, pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Masks the pred and target tensors according to nnunet ``ignore_label``. The number of classes in the input @@ -748,8 +759,13 @@ 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("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): From 34f35c72a6a0407a97a710ac8ed0ec4606906298 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 21 May 2025 13:28:17 -0400 Subject: [PATCH 05/32] working --- examples/nnunet_pfl_example/README.md | 69 +++++ examples/nnunet_pfl_example/__init__.py | 0 examples/nnunet_pfl_example/client.py | 245 ++++++++++++++++++ examples/nnunet_pfl_example/config.yaml | 8 + examples/nnunet_pfl_example/server.py | 184 +++++++++++++ .../fedavg_with_adaptive_constraint.py | 7 +- 6 files changed, 510 insertions(+), 3 deletions(-) create mode 100644 examples/nnunet_pfl_example/README.md create mode 100644 examples/nnunet_pfl_example/__init__.py create mode 100644 examples/nnunet_pfl_example/client.py create mode 100644 examples/nnunet_pfl_example/config.yaml create mode 100644 examples/nnunet_pfl_example/server.py diff --git a/examples/nnunet_pfl_example/README.md b/examples/nnunet_pfl_example/README.md new file mode 100644 index 000000000..e5f890955 --- /dev/null +++ b/examples/nnunet_pfl_example/README.md @@ -0,0 +1,69 @@ +# NnUNetClient Example + +This example demonstrates how to use the NnunetClient to train nnunet segmentation models in a federated setting. + +By default this example trains an nnunet model on the Task04_Hippocampus dataset from the Medical Segmentation Decathlon (MSD). However, any of the MSD datasets can be used by specifying them with the msd_dataset_name flag for the client. To run this example first create a config file for the server. An example config has been provided in this directory. The required keys for the config are: + +```yaml +# Parameters that describe the server +n_server_rounds: 1 + +# Parameters that describe the clients +n_clients: 1 +local_epochs: 1 # Or local_steps, one or the other must be chosen + +nnunet_config: 2d +``` + +The only additional parameter required by nnunet is nnunet_config which is one of the official nnunet configurations (2d, 3d_fullres, 3d_lowres, 3d_cascade_fullres) + +One may also add the following optional keys to the config yaml file. If a nnunet plans file (which specifies model architecture and training hyperparameters) is not provided, the server will ask one of the clients to generate one using nnunet. + +```yaml +# Optional config parameters +nnunet_plans: /Path/to/nnunet/plans.json +starting_checkpoint: /Path/to/starting/checkpoint.pth +``` + +To run a federated learning experiment with nnunet models, first ensure you are in the FL4Health directory and then start the nnunet server using the following command. To view a list of optional flags use the --help flag + +```bash +python -m examples.nnunet_example.server --config_path examples/nnunet_example/config.yaml +``` + +Once the server has started, start the necessary number of clients specified by the n_clients key in the config file. Each client can be started by running the following command in a separate session. To view a list of optional flags use the --help flag. + +```bash +python -m examples.nnunet_example.client --dataset_path examples/datasets/nnunet +``` + +The MSD dataset will be downloaded and prepared automatically by the nnunet example script if it does not already exist. The dataset_path flag is used as more of a data working directory by the client. The client will create nnunet_raw, nnunet_preprocessed and nnunet_results sub directories if they do not already exist in the dataset_path folder. The dataset itself will be stored in a folder within nnunet_raw. Therefore when checking if the data already exists, the client will look for the following folder '{dataset_path}/nnunet_raw/{dataset_name}' + +## Definitions + +### Logits: + +The outputs of a model prior to the activation function. Values are unconstrained (-inf, inf) + +### Probabilities: + +Probabilities are values constrained to the range (0, 1) that represent the models confidence of a pixel/voxel being part of a particular class. Similar to other DNN's, they do not necessarily represent actual probabilities, however it is sometimes convenient to interpret them as such. The predicted probabilities are the outputs of a model after a normalizing activation function such as softmax or sigmoid. A 2d example is shown below. In some instances one might have labels that are probabilities as opposed to integer class labels in which case we would refer to them as ground-truth probabilities. + +### Detection Maps: + +Images that contain an arbitrary number of distinct detected volumes generally derived from the predicted probabilities of a segmentation model. Values are constrained to range [0, 1]. Detected volumes are defined as: +- Each detected volume is a connected component that must be non-connected and non-overlapping (mutually exclusive) with other volumes of the same class. (Therefore detection maps for multiclass segmentation must be one hot encoded) +- Each pixel/voxel within a volume must have the same predicted probability. Therefore there is a single confidence/likelihood score for each volume. + +Detected volumes typically also have a minimum size determined by the number of pixels/voxels that are a part of the volume. Detection maps may be computed from probabilities in a variety of ways. One example used for 3d medical images can be found in the [report guided annotation](https://github.com/DIAGNijmegen/Report-Guided-Annotation) API. An example of a 2d detection map is shown below. + +### Segmentations: + +Images in which pixels have been labelled or assigned one or more specific integer classes. If one hot encoded they must be binary {0, 1} or boolean {False, True} tensors. If not one hot encoded they must be be tensors containing only integers that represent the class labels (eg. constrained to {0, 1, 2, ..., N}). The labels/targets for segmentation models may be referred to as ground-truth segmentations. Predicted segmentations refers to model outputs (which are usually probability maps) that have been processed or thresholded in some way to adhere to the definition of a segmentation. They usually represents the model's final prediction of the class with no information on confidence. An example of a binary or one-hot-encoded predicted segmentation is shown below. + +#### Examples of different output types in the case of 2d binary segmentation: + + + | | | +:----------------------------:|:------------------------------:|:----------------------------------: +Probabilities | Detection Map | Segmentation| diff --git a/examples/nnunet_pfl_example/__init__.py b/examples/nnunet_pfl_example/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/examples/nnunet_pfl_example/client.py b/examples/nnunet_pfl_example/client.py new file mode 100644 index 000000000..18ee2359a --- /dev/null +++ b/examples/nnunet_pfl_example/client.py @@ -0,0 +1,245 @@ +import argparse +import os +import warnings +from logging import DEBUG, INFO +from os.path import exists, join +from pathlib import Path +from typing import Any, Literal + +from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer +from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule + +with warnings.catch_warnings(): + # Silence deprecation warnings from sentry sdk due to flwr and wandb + # https://github.com/adap/flower/issues/4086 + warnings.filterwarnings("ignore", category=DeprecationWarning) + import wandb # noqa: F401 + +import torch +from flwr.client import start_client +from flwr.common.logger import log, update_console_handler +from nnunetv2.dataset_conversion.convert_MSD_dataset import convert_msd_dataset +from torchmetrics.segmentation import GeneralizedDiceScore + +from fl4health.clients.nnunet_client import NnunetClient +from fl4health.metrics import TorchMetric +from fl4health.metrics.compound_metrics import TransformsMetric +from fl4health.mixins.personalized import PersonalizedMode, make_it_personal +from fl4health.utils.load_data import load_msd_dataset +from fl4health.utils.msd_dataset_sources import get_msd_dataset_enum, msd_num_labels +from fl4health.utils.nnunet_utils import get_segs_from_probs, set_nnunet_env + +personalized_client_classes = {"ditto": make_it_personal(NnunetClient, PersonalizedMode.DITTO)} + + +def main( + dataset_path: Path, + msd_dataset_name: str, + server_address: str, + fold: int | str, + always_preprocess: bool = False, + verbose: bool = True, + compile: bool = True, + intermediate_client_state_dir: str | None = None, + client_name: str | None = None, + personalized_strategy: Literal["ditto"] | None = None, +) -> None: + with torch.autograd.set_detect_anomaly(True): + # Log device and server address + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + log(INFO, f"Using device: {device}") + log(INFO, f"Using server address: {server_address}") + + # Load the dataset if necessary + msd_dataset_enum = get_msd_dataset_enum(msd_dataset_name) + nnUNet_raw = join(dataset_path, "nnunet_raw") + if not exists(join(nnUNet_raw, msd_dataset_enum.value)): + log(INFO, f"Downloading and extracting {msd_dataset_enum.value} dataset") + load_msd_dataset(nnUNet_raw, msd_dataset_name) + + # The dataset ID will be the same as the MSD Task number + dataset_id = int(msd_dataset_enum.value[4:6]) + nnunet_dataset_name = f"Dataset{dataset_id:03d}_{msd_dataset_enum.value.split('_')[1]}" + + # Convert the msd dataset if necessary + if not exists(join(nnUNet_raw, nnunet_dataset_name)): + log(INFO, f"Converting {msd_dataset_enum.value} into nnunet dataset") + convert_msd_dataset(source_folder=join(nnUNet_raw, msd_dataset_enum.value)) + + # Create a metric + dice = TransformsMetric( + metric=TorchMetric( + name="Pseudo DICE", + metric=GeneralizedDiceScore( + num_classes=msd_num_labels[msd_dataset_enum], weight_type="square", include_background=False + ).to(device), + ), + pred_transforms=[torch.sigmoid, get_segs_from_probs], + ) + + if intermediate_client_state_dir is not None: + checkpoint_and_state_module = ClientCheckpointAndStateModule( + state_checkpointer=PerRoundStateCheckpointer(Path(intermediate_client_state_dir)) + ) + else: + checkpoint_and_state_module = None + + # Create client + client_kwargs: dict[str, Any] = {} + client_kwargs.update( + # Args specific to nnUNetClient + dataset_id=dataset_id, + fold=fold, + always_preprocess=always_preprocess, + verbose=verbose, + compile=compile, + # BaseClient Args + device=device, + metrics=[dice], + progress_bar=verbose, + checkpoint_and_state_module=checkpoint_and_state_module, + client_name=client_name, + ) + if personalized_strategy: + log(INFO, f"Setting up client for personalized strategy: {personalized_strategy}") + client = personalized_client_classes[personalized_strategy](**client_kwargs) + else: + log(INFO, "Setting up client without personalization") + client = NnunetClient(**client_kwargs) + log(INFO, f"Using client: {type(client).__name__}") + + start_client(server_address=server_address, client=client.to_client()) + + # Shutdown the client + client.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="nnunet_example/client.py", + description="An exampled of nnUNetClient on any of the Medical \ + Segmentation Decathlon (MSD) datasets. Automatically generates a \ + nnunet segmentation model and trains it in a federated setting", + ) + + # I have to use underscores instead of dashes because thats how they + # defined it in run_smoke_tests + parser.add_argument( + "--dataset_path", + type=str, + required=True, + help="Path to the folder in which data should be stored. This script \ + will automatically create nnunet_raw, and nnunet_preprocessed \ + subfolders if they don't already exist. This script will also \ + attempt to download and prepare the MSD Dataset into the \ + nnunet_raw folder if it does not already exist.", + ) + parser.add_argument( + "--fold", + type=str, + required=False, + default="0", + help="[OPTIONAL] Which fold of the local client dataset to use for \ + validation. nnunet defaults to 5 folds (0 to 4). Can also be set \ + to 'all' to use all the data for both training and validation. \ + Defaults to fold 0", + ) + parser.add_argument( + "--msd_dataset_name", + type=str, + required=False, + default="Task04_Hippocampus", # The smallest dataset + help="[OPTIONAL] Name of the MSD dataset to use. The options are \ + defined by the values of the MsdDataset enum as returned by the \ + get_msd_dataset_enum function", + ) + parser.add_argument( + "--always-preprocess", + action="store_true", + required=False, + help="[OPTIONAL] Use this to force preprocessing the nnunet data \ + even if the preprocessed data is found to already exist", + ) + parser.add_argument( + "--server_address", + type=str, + required=False, + default="0.0.0.0:8080", + help="[OPTIONAL] The server address for the clients to communicate \ + to the server through. Defaults to 0.0.0.0:8080", + ) + parser.add_argument( + "--verbose", + action="store_true", + required=False, + help="[OPTIONAL] Use this flag to see extra INFO logs and a progress bar", + ) + parser.add_argument( + "--debug", + help="[OPTIONAL] Include flag to print DEBUG logs", + action="store_const", + dest="logLevel", + const=DEBUG, + default=INFO, + ) + parser.add_argument( + "--skip-compile", + action="store_true", + required=False, + help="[OPTIONAL] Include flag to train without jit compiling the pytorch model first", + ) + + parser.add_argument( + "--intermediate-client-state-dir", + type=str, + required=False, + default=None, + help="[OPTIONAL] Directory to store client state during training. Defaults to None", + ) + parser.add_argument( + "--client-name", + type=str, + required=False, + default=None, + help="[OPTIONAL] Name of the client used to name client state checkpoint. \ + Defaults to None, in which case a random name is generated for the client", + ) + parser.add_argument( + "--personalized-strategy", + type=str, + required=False, + default=None, + help="[OPTIONAL] Personalized strategy to use. Can be 'ditto' or 'mr-mtl' \ + Defaults to None, in which no personalized strategy is applied.", + ) + + args = parser.parse_args() + + # Set the log level + update_console_handler(level=args.logLevel) + + # Create nnunet directory structure and set environment variables + nnUNet_raw = join(args.dataset_path, "nnunet_raw") + nnUNet_preprocessed = join(args.dataset_path, "nnunet_preprocessed") + os.makedirs(nnUNet_raw, exist_ok=True) + os.makedirs(nnUNet_preprocessed, exist_ok=True) + set_nnunet_env( + nnUNet_raw=nnUNet_raw, + nnUNet_preprocessed=nnUNet_preprocessed, + nnUNet_results=join(args.dataset_path, "nnUNet_results"), + ) + + # Check fold argument and start main method + fold: int | str = "all" if args.fold == "all" else int(args.fold) + main( + dataset_path=Path(args.dataset_path), + msd_dataset_name=args.msd_dataset_name, + server_address=args.server_address, + fold=fold, + always_preprocess=args.always_preprocess, + verbose=args.verbose, + compile=not args.skip_compile, + intermediate_client_state_dir=args.intermediate_client_state_dir, + client_name=args.client_name, + personalized_strategy=args.personalized_strategy, + ) diff --git a/examples/nnunet_pfl_example/config.yaml b/examples/nnunet_pfl_example/config.yaml new file mode 100644 index 000000000..d2371d0d6 --- /dev/null +++ b/examples/nnunet_pfl_example/config.yaml @@ -0,0 +1,8 @@ +# Parameters that describe the server +n_server_rounds: 1 + +# Parameters that describe the clients +n_clients: 1 +local_epochs: 1 + +nnunet_config: 2d diff --git a/examples/nnunet_pfl_example/server.py b/examples/nnunet_pfl_example/server.py new file mode 100644 index 000000000..1cfe9fa8f --- /dev/null +++ b/examples/nnunet_pfl_example/server.py @@ -0,0 +1,184 @@ +import argparse +import json +import pickle +from functools import partial +from pathlib import Path + +import flwr as fl +import torch +import yaml +from flwr.common.parameter import ndarrays_to_parameters +from flwr.common.typing import Config +from flwr.server.client_manager import SimpleClientManager + +from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer +from fl4health.checkpointing.server_module import NnUnetServerCheckpointAndStateModule +from fl4health.metrics.metric_aggregation import evaluate_metrics_aggregation_fn, fit_metrics_aggregation_fn +from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger +from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking +from fl4health.parameter_exchange.parameter_packer import ( + ParameterPackerAdaptiveConstraint, +) +from fl4health.servers.nnunet_server import NnunetServer +from fl4health.strategies.fedavg_with_adaptive_constraint import FedAvgWithAdaptiveConstraint +from fl4health.utils.config import make_dict_with_epochs_or_steps +from fl4health.utils.random import set_all_random_seeds + + +def get_config( + current_server_round: int, + nnunet_config: str, + n_server_rounds: int, + batch_size: int, + n_clients: int, + nnunet_plans: str | None = None, + local_epochs: int | None = None, + local_steps: int | None = None, +) -> Config: + # Create config + config: Config = { + "n_clients": n_clients, + "nnunet_config": nnunet_config, + "n_server_rounds": n_server_rounds, + "batch_size": batch_size, + **make_dict_with_epochs_or_steps(local_epochs, local_steps), + "current_server_round": current_server_round, + } + + # Check if plans were provided + if nnunet_plans is not None: + plans_bytes = pickle.dumps(json.load(open(nnunet_plans, "r"))) + config["nnunet_plans"] = plans_bytes + + return config + + +def main( + config: dict, + server_address: str, + intermediate_server_state_dir: str | None = None, + server_name: str | None = None, +) -> None: + # Partial function with everything set except current server round + fit_config_fn = partial( + get_config, + n_clients=config["n_clients"], + nnunet_config=config["nnunet_config"], + n_server_rounds=config["n_server_rounds"], + batch_size=0, # Set this to 0 because we're not using it + nnunet_plans=config.get("nnunet_plans"), + local_epochs=config.get("local_epochs"), + local_steps=config.get("local_steps"), + ) + + if config.get("starting_checkpoint"): + model = torch.load(config["starting_checkpoint"]) + # Of course nnunet stores their pytorch models differently. + params = ndarrays_to_parameters([val.cpu().numpy() for _, val in model["network_weights"].items()]) + else: + params = None + + strategy = FedAvgWithAdaptiveConstraint( + min_fit_clients=config["n_clients"], + min_evaluate_clients=config["n_clients"], + # Server waits for min_available_clients before starting FL rounds + min_available_clients=config["n_clients"], + on_fit_config_fn=fit_config_fn, + # We use the same fit config function, as nothing changes for eval + on_evaluate_config_fn=fit_config_fn, + fit_metrics_aggregation_fn=fit_metrics_aggregation_fn, + evaluate_metrics_aggregation_fn=evaluate_metrics_aggregation_fn, + initial_parameters=params, + initial_loss_weight=0.1, + adapt_loss_weight=False, + ) + + state_checkpointer = ( + PerRoundStateCheckpointer(Path(intermediate_server_state_dir)) + if intermediate_server_state_dir is not None + else None + ) + checkpoint_and_state_module = NnUnetServerCheckpointAndStateModule( + model=None, parameter_exchanger=FullParameterExchanger(), state_checkpointer=state_checkpointer + ) + + checkpoint_and_state_module.parameter_exchanger = FullParameterExchangerWithPacking( # type:ignore [assignment] + ParameterPackerAdaptiveConstraint() + ) + + server = NnunetServer( + client_manager=SimpleClientManager(), + fl_config=config, + # The fit_config_fn contains all of the necessary information for param initialization, so we reuse it here + on_init_parameters_config_fn=fit_config_fn, + strategy=strategy, + checkpoint_and_state_module=checkpoint_and_state_module, + server_name=server_name, + accept_failures=False, + ) + + fl.server.start_server( + server=server, + server_address=server_address, + config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]), + ) + + # Shutdown server + server.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--config_path", + action="store", + type=str, + help="Path to the configuration file. See examples/nnunet_example/README.md for more info", + ) + parser.add_argument( + "--server-address", + type=str, + required=False, + default="0.0.0.0:8080", + help="[OPTIONAL] The address to use for the server. Defaults to \ + 0.0.0.0:8080", + ) + + parser.add_argument( + "--intermediate-server-state-dir", + type=str, + required=False, + default=None, + help="[OPTIONAL] Directory to store server state. Defaults to None", + ) + parser.add_argument( + "--server-name", + type=str, + required=False, + default=None, + help="[OPTIONAL] Name of the server used as file name when \ + checkpointing server state. Defaults to \ + None, in which case the server will generate random name \ + ", + ) + parser.add_argument( + "--seed", + action="store", + type=int, + help="Seed for the random number generators across python, torch, and numpy", + required=False, + ) + args = parser.parse_args() + + with open(args.config_path, "r") as f: + config = yaml.safe_load(f) + + # Set the random seed for reproducibility + set_all_random_seeds(args.seed) + + main( + config, + server_address=args.server_address, + intermediate_server_state_dir=args.intermediate_server_state_dir, + server_name=args.server_name, + ) diff --git a/fl4health/strategies/fedavg_with_adaptive_constraint.py b/fl4health/strategies/fedavg_with_adaptive_constraint.py index 28f0f1370..2fb1227fb 100644 --- a/fl4health/strategies/fedavg_with_adaptive_constraint.py +++ b/fl4health/strategies/fedavg_with_adaptive_constraint.py @@ -28,7 +28,7 @@ def __init__( on_fit_config_fn: Callable[[int], dict[str, Scalar]] | None = None, on_evaluate_config_fn: Callable[[int], dict[str, Scalar]] | None = None, accept_failures: bool = True, - initial_parameters: Parameters, + initial_parameters: Parameters | None, fit_metrics_aggregation_fn: MetricsAggregationFn | None = None, evaluate_metrics_aggregation_fn: MetricsAggregationFn | None = None, initial_loss_weight: float = 1.0, @@ -102,8 +102,9 @@ def __init__( self.previous_loss = float("inf") - self.server_model_weights = parameters_to_ndarrays(initial_parameters) - initial_parameters.tensors.extend(ndarrays_to_parameters([np.array(initial_loss_weight)]).tensors) + if initial_parameters: + self.server_model_weights = parameters_to_ndarrays(initial_parameters) + initial_parameters.tensors.extend(ndarrays_to_parameters([np.array(initial_loss_weight)]).tensors) super().__init__( fraction_fit=fraction_fit, From 407858fe975910fa1d96a8b35e79e4524fc47a65 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Thu, 22 May 2025 12:17:15 -0400 Subject: [PATCH 06/32] update mixins --- fl4health/clients/nnunet_client.py | 66 ++++--- .../mixins/adaptive_drift_constrained.py | 69 +++---- fl4health/mixins/personalized/ditto.py | 176 ++++-------------- 3 files changed, 107 insertions(+), 204 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 9d92444f9..48cbd0a31 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -206,33 +206,19 @@ def __init__( if self.verbose: log(INFO, "Disabling model optimizations and JIT compilation. This may impact runtime performance.") - def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: - """ - Given a single batch of input and target data, generate predictions, compute loss, update parameters and - optionally update metrics if they exist. (i.e. backprop on a single batch of data). Assumes ``self.model`` is - in train mode already. - - Overrides parent to include mixed precision training (autocasting and corresponding gradient scaling) as per - the original ``nnUNetTrainer``. - - Args: - input (TorchInputType): The input to be fed into the model. - target (TorchTargetType): The target corresponding to the input. - - Returns: - Tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any - predictions produced by the model. - """ + def _train_step( + self, model: nn.Module, optimizer: ..., input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": return super().train_step(input, target) # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler # Clear gradients from optimizer if they exist - self.optimizers["local"].zero_grad() + optimizer.zero_grad() # Call user defined methods to get predictions and compute loss - preds, features = self.predict(input) + preds, features = self._predict(model, input) target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) @@ -241,15 +227,35 @@ 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["local"]) - self.transform_gradients(losses) + self.grad_scaler.unscale_(optimizer) + self._transform_gradients(model, losses) # Update parameters and scaler - self.grad_scaler.step(self.optimizers["local"]) + self.grad_scaler.step(optimizer) self.grad_scaler.update() return losses, preds + def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: + """ + Given a single batch of input and target data, generate predictions, compute loss, update parameters and + optionally update metrics if they exist. (i.e. backprop on a single batch of data). Assumes ``self.model`` is + in train mode already. + + Overrides parent to include mixed precision training (autocasting and corresponding gradient scaling) as per + the original ``nnUNetTrainer``. + + Args: + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + Tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any + predictions produced by the model. + """ + # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. + return self._train_step(self.model, self.optimizers["global"], input, target) + @use_default_signal_handlers # Dataloaders use multiprocessing def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]: """ @@ -759,13 +765,16 @@ 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} + + log(DEBUG, f"preds: {preds.keys()}") + + # If personalized, the + preds = {k: v for k, v in preds.items() if not k.startswith("global")} # remove prefix preds = {k.replace("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): @@ -822,7 +831,7 @@ def get_client_specific_logs( logging_mode: LoggingMode, ) -> tuple[str, list[tuple[LogLevel, str]]]: if logging_mode == LoggingMode.TRAIN: - lr = float(self.optimizers["local"].param_groups[0]["lr"]) + lr = float(self.optimizers["global"].param_groups[0]["lr"]) if current_epoch is None: # Assume training by steps return f"Initial LR {lr}", [] @@ -832,7 +841,7 @@ def get_client_specific_logs( return "", [] def get_client_specific_reports(self) -> dict[str, Any]: - return {"learning_rate": float(self.optimizers["local"].param_groups[0]["lr"])} + return {"learning_rate": float(self.optimizers["global"].param_groups[0]["lr"])} @use_default_signal_handlers # Experiment planner spawns a process I think def get_properties(self, config: Config) -> dict[str, Scalar]: @@ -936,6 +945,9 @@ 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, model: torch.nn.Module, losses: TrainingLosses) -> None: + nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm) + def transform_gradients(self, losses: TrainingLosses) -> None: """ Apply the gradient clipping performed by the default nnunet trainer. This is the default behavior for @@ -944,4 +956,4 @@ def transform_gradients(self, losses: TrainingLosses) -> None: Args: losses (TrainingLosses): Not used for this transformation. """ - nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm) + self._transform_gradients(self.model, self.max_grad_norm) diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py index e907450f4..841bd0d7e 100644 --- a/fl4health/mixins/adaptive_drift_constrained.py +++ b/fl4health/mixins/adaptive_drift_constrained.py @@ -1,7 +1,7 @@ """AdaptiveDriftConstrainedMixin""" import warnings -from logging import INFO, WARN +from logging import INFO, WARN, DEBUG from typing import Any, Protocol, runtime_checkable import torch @@ -148,39 +148,40 @@ def set_parameters( super().set_parameters(server_model_state, config, fitting_round) # type: ignore[safe-super] - def compute_training_loss( - self: AdaptiveDriftConstrainedProtocol, - preds: TorchPredType, - features: TorchFeatureType, - target: TorchTargetType, - ) -> TrainingLosses: - """ - Computes training loss given predictions of the model and ground truth data. Adds to objective by including - penalty loss. - - Args: - preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in - dictionary will be used to compute metrics. - features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. - target: (TorchTargetType): Ground truth data to evaluate predictions against. - - Returns: - TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed - by name. Additional losses includes penalty loss. - """ - loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) - if additional_losses is None: - additional_losses = {} - - additional_losses["loss"] = loss.clone() - # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation - additional_losses["loss_for_adaptation"] = loss.clone() - - # Compute the drift penalty loss and store it in the additional losses dictionary. - penalty_loss = self.compute_penalty_loss() - additional_losses["penalty_loss"] = penalty_loss.clone() - - return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) + # def compute_training_loss( + # self: AdaptiveDriftConstrainedProtocol, + # preds: TorchPredType, + # features: TorchFeatureType, + # target: TorchTargetType, + # ) -> TrainingLosses: + # """ + # Computes training loss given predictions of the model and ground truth data. Adds to objective by including + # penalty loss. + + # Args: + # preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in + # dictionary will be used to compute metrics. + # features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. + # target: (TorchTargetType): Ground truth data to evaluate predictions against. + + # Returns: + # TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed + # by name. Additional losses includes penalty loss. + # """ + # log(DEBUG, "COMPUTE TRAINING LOSS CALL FROM ADAPTIVE MIXIN") + # loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) + # if additional_losses is None: + # additional_losses = {} + + # additional_losses["loss"] = loss.clone() + # # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation + # additional_losses["loss_for_adaptation"] = loss.clone() + + # # Compute the drift penalty loss and store it in the additional losses dictionary. + # penalty_loss = self.compute_penalty_loss() + # additional_losses["penalty_loss"] = penalty_loss.clone() + + # return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) def get_parameter_exchanger(self: AdaptiveDriftConstrainedProtocol, config: Config) -> ParameterExchanger: """ diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index a2d9df064..2b21e522e 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -74,6 +74,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: if not isinstance(self, BasicClientProtocolPreSetup): raise RuntimeError("This object needs to satisfy `BasicClientProtocolPreSetup`.") + # + def __init_subclass__(cls, **kwargs: Any) -> None: """This method is called when a class inherits from AdaptiveMixin""" super().__init_subclass__(**kwargs) @@ -370,73 +372,48 @@ def train_step( corresponding to predictions from the global and local Ditto models for metric evaluations. """ - # Clear gradients from optimizers if they exist - self.optimizers["global"].zero_grad() - self.optimizers["local"].zero_grad() - - # Forward pass on both the global and local models - preds, features = self.predict(input) - target = self.transform_target(target) # Apply transformation (Defaults to identity) + # global + global_losses, global_preds = self._train_step(self.global_model, self.optimizers["global"], input, target) + # local + local_losses, local_preds = self._train_step(self.model, self.optimizers["local"], input, target) - # Compute all relevant losses - losses = self.compute_training_loss(preds, features, target) + log(DEBUG, f"global_losses: {type(global_losses)}") + log(DEBUG, f"local_losses: {type(local_losses)}") - # Take a step with the global model vanilla loss - losses.additional_losses["global_loss"].backward() - self.optimizers["global"].step() + local_losses = cast(TrainingLosses, local_losses) + global_losses = cast(TrainingLosses, global_losses) - # Take a step with the local model using the local loss and Ditto constraint - losses.backward["backward"].backward() - self.optimizers["local"].step() + # aggregate global and local losses - # Return dictionary of predictions where key is used to name respective MetricMeters - return losses, preds - - def predict( - self: DittoPersonalizedProtocol, - input: TorchInputType, - ) -> tuple[TorchPredType, TorchFeatureType]: - """ - Computes the predictions for both the **GLOBAL** and **LOCAL** models and pack them into the prediction - dictionary + local_loss = local_losses.backward["backward"] + global_loss = global_losses.backward["backward"] - Args: - input (TorchInputType): Inputs to be fed into both models. + # from compute_loss_and_additional_losses + additional_losses = { + "local_loss": local_loss.clone(), + "global_loss": global_loss, + } - Returns: - tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains predictions indexed by - name and the second element contains intermediate activations index by name. For Ditto, we only need the - predictions, so the second dictionary is simply empty. + # from compute_loss + # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or + # less weight is used to constrain it to the global model (as in FedProx) + additional_losses["loss_for_adaptation"] = additional_losses["local_loss"].clone() - Raises: - ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model - forward. - """ + # This is the Ditto penalty loss of the local model compared with the original Global model weights, scaled + # by drift_penalty_weight (or lambda in the original paper) + penalty_loss = self.compute_penalty_loss() + additional_losses["penalty_loss"] = penalty_loss.clone() - if hasattr(self, "_predict"): - log(INFO, "Using '_predict' to make predictions") - global_preds, _ = self._predict(self.safe_global_model(), input) - local_preds, _ = self._predict(self.model, input) - log(INFO, "Successfully predicted for global and local models") - else: - if isinstance(input, torch.Tensor): - global_preds = self.safe_global_model()(input) - local_preds = self.model(input) - elif isinstance(input, dict): - # If input is a dictionary, then we unpack it before computing the forward pass. - # Note that this assumes the keys of the input match (exactly) the keyword args - # of the forward method. - global_preds = self.safe_global_model()(**input) - local_preds = self.model(**input) + training_losses = TrainingLosses(backward=local_loss + penalty_loss, additional_losses=additional_losses) + # aggregate preds if isinstance(global_preds, torch.Tensor) and isinstance(local_preds, torch.Tensor): - return {"global": global_preds, "local": local_preds}, {} + agg_preds = {"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 types. Global: {type(global_preds)}, local: {type(local_preds)}") + agg_preds = {f"global-{k}": v for k, v in global_preds.items()} + agg_preds.update(**{f"local-{k}": v for k, v in local_preds.items()}) + + return training_losses, agg_preds def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Helper method to extract predictions from global and local models. @@ -460,93 +437,6 @@ def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, retval = {k.replace(f"{kind}-", ""): v for k, v in retval.items()} return retval - def compute_loss_and_additional_losses( - self: DittoPersonalizedProtocol, - preds: TorchPredType, - features: TorchFeatureType, - target: TorchTargetType, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: - """ - Computes the local model loss and the global Ditto model loss (stored in additional losses) for reporting and - training of the global model - - Args: - preds (TorchPredType): Prediction(s) of the model(s) indexed by name. - features (TorchFeatureType): Feature(s) of the model(s) indexed by name. - target (TorchTargetType): Ground truth data to evaluate predictions against. - Returns: - tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple with: - - - The tensor for the model loss - - 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) - - log(DEBUG, f"global_preds has keys {global_preds.keys()}") - log(DEBUG, f"local_preds has keys {local_preds.keys()}") - - # 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(global_preds, features, target) - - # Compute local model loss - local_loss, _ = self._special_compute_loss_and_additional_losses(local_preds, features, target) - - else: - global_loss = self.criterion(global_preds["global"], target) - - # Compute local model loss - local_loss = self.criterion(local_preds["local"], target) - - additional_losses = {"local_loss": local_loss.clone(), "global_loss": global_loss} - - return local_loss, additional_losses - - def compute_training_loss( - self: DittoPersonalizedProtocol, - preds: TorchPredType, - features: TorchFeatureType, - target: TorchTargetType, - ) -> TrainingLosses: - """ - Computes training losses given predictions of the global and local models and ground truth data. - For the local model we add to the vanilla loss function by including Ditto penalty loss which is the l2 inner - product between the initial global model weights and weights of the local model. This is stored in backward - The loss to optimize the global model is stored in the additional losses dictionary under "global_loss" - - Args: - preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in - dictionary will be used to compute metrics. - features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. - target: (TorchTargetType): Ground truth data to evaluate predictions against. - - Returns: - TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed by - name. Additional losses includes each loss component and the global model - loss tensor. - """ - # Check that both models are in training mode - assert self.safe_global_model().training and self.model.training - - # local loss is stored in loss, global model loss is stored in additional losses. - loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) - additional_losses = additional_losses or {} # make mypy happy - - # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or - # less weight is used to constrain it to the global model (as in FedProx) - additional_losses["loss_for_adaptation"] = additional_losses["local_loss"].clone() - - # This is the Ditto penalty loss of the local model compared with the original Global model weights, scaled - # by drift_penalty_weight (or lambda in the original paper) - penalty_loss = self.compute_penalty_loss() - additional_losses["penalty_loss"] = penalty_loss.clone() - - return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) - @ensure_protocol_compliance def validate( self: DittoPersonalizedProtocol, include_losses_in_metrics: bool = False From 2ba4f7dffb8131cb521b3d855062dd860381fd3c Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Mon, 26 May 2025 10:46:39 -0400 Subject: [PATCH 07/32] update basic client with helper methods --- fl4health/clients/basic_client.py | 100 ++++++++++++++++--------- fl4health/clients/nnunet_client.py | 4 +- fl4health/mixins/core_protocols.py | 5 ++ fl4health/mixins/personalized/ditto.py | 4 +- 4 files changed, 74 insertions(+), 39 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index eaff0df1d..1d0f01e06 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -559,25 +559,20 @@ def update_metric_manager( """ metric_manager.update(preds, target) - def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: - """ - Given a single batch of input and target data, generate predictions, compute loss, update parameters and - optionally update metrics if they exist. (i.e. backprop on a single batch of data). - Assumes ``self.model`` is in train mode already. + def _train_step( + self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + """Helper train step. - Args: - input (TorchInputType): The input to be fed into the model. - target (TorchTargetType): The target corresponding to the input. - - Returns: - tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with - a dictionary of any predictions produced by the model. + This interface allows for injection of model and optimizer params, which + are useful for personalized FL methods. """ + # Clear gradients from optimizer if they exist - self.optimizers["global"].zero_grad() + optimizer.zero_grad() # Call user defined methods to get predictions and compute loss - preds, features = self.predict(input) + preds, features = self._predict(model, input) target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) @@ -588,6 +583,22 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr return losses, preds + def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: + """ + Given a single batch of input and target data, generate predictions, compute loss, update parameters and + optionally update metrics if they exist. (i.e. backprop on a single batch of data). + Assumes ``self.model`` is in train mode already. + + Args: + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with + a dictionary of any predictions produced by the model. + """ + return self._train_step(self.model, self.optimizers["global"], input, target) + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: """ Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already. @@ -975,34 +986,19 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger: """ return FullParameterExchanger() - def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: - """ - Computes the prediction(s), and potentially features, of the model(s) given the input. - - Args: - input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``, - it is assumed that the keys of input match the names of the keyword arguments of - ``self.model.forward().`` - - Returns: - tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of - predictions indexed by name and the second element contains intermediate activations indexed by name. By - passing features, we can compute losses such as the contrastive loss in MOON. All predictions included in - dictionary will by default be used to compute metrics separately. + def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + """Helper predict method. - Raises: - TypeError: Occurs when something other than a tensor or dict of tensors is passed in to the model's - forward method. - ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model - forward. + Unlike, predict(), this interface allows for injecting the model param. """ + if isinstance(input, torch.Tensor): - output = self.model(input) + output = model(input) elif isinstance(input, dict): # If input is a dictionary, then we unpack it before computing the forward pass. # Note that this assumes the keys of the input match (exactly) the keyword args # of self.model.forward(). - output = self.model(**input) + output = model(**input) else: raise TypeError("'input' must be of type torch.Tensor or dict[str, torch.Tensor].") @@ -1018,6 +1014,29 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp else: raise ValueError("Model forward did not return a tensor, dictionary of tensors, or tuple of tensors") + def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + """ + Computes the prediction(s), and potentially features, of the model(s) given the input. + + Args: + input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``, + it is assumed that the keys of input match the names of the keyword arguments of + ``self.model.forward().`` + + Returns: + tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of + predictions indexed by name and the second element contains intermediate activations indexed by name. By + passing features, we can compute losses such as the contrastive loss in MOON. All predictions included in + dictionary will by default be used to compute metrics separately. + + Raises: + TypeError: Occurs when something other than a tensor or dict of tensors is passed in to the model's + forward method. + ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model + forward. + """ + return self._predict(self.model, input) + def compute_loss_and_additional_losses( self, preds: TorchPredType, features: TorchFeatureType, target: TorchTargetType ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: @@ -1280,6 +1299,15 @@ def update_before_epoch(self, epoch: int) -> None: """ pass + def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: + """ + Helper transform gradients method. + + Unlike transform_gradients(), this helper's interface allows for injecting + model as a param. + """ + pass + def transform_gradients(self, losses: TrainingLosses) -> None: """ Hook function for model training only called after backwards pass but before optimizer step. Useful for @@ -1288,7 +1316,7 @@ def transform_gradients(self, losses: TrainingLosses) -> None: Args: losses (TrainingLosses): The losses object from the train step """ - pass + return self._transform_gradients(self.model, losses) def _save_client_state(self) -> None: """ diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 48cbd0a31..322f2d8c0 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -207,7 +207,7 @@ def __init__( log(INFO, "Disabling model optimizations and JIT compilation. This may impact runtime performance.") def _train_step( - self, model: nn.Module, optimizer: ..., input: TorchInputType, target: TorchTargetType + self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": @@ -956,4 +956,4 @@ def transform_gradients(self, losses: TrainingLosses) -> None: Args: losses (TrainingLosses): Not used for this transformation. """ - self._transform_gradients(self.model, self.max_grad_norm) + self._transform_gradients(self.model, losses) diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index d28ff7794..df7c62a29 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -77,6 +77,11 @@ def initialize_all_model_weights(self, parameters: NDArrays, config: Config) -> def update_before_train(self, current_server_round: int) -> None: pass # pragma: no cover + def _train_step( + self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + pass # pragma: no cover + def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: pass # pragma: no cover diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index 2b21e522e..94c25d025 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -373,7 +373,9 @@ def train_step( """ # global - global_losses, global_preds = self._train_step(self.global_model, self.optimizers["global"], input, target) + global_losses, global_preds = self._train_step( + self.safe_global_model(), self.optimizers["global"], input, target + ) # local local_losses, local_preds = self._train_step(self.model, self.optimizers["local"], input, target) From d5e139de1b4c91a287649e972fe3e122eae1bbf7 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Mon, 26 May 2025 11:58:59 -0400 Subject: [PATCH 08/32] update adaptive to override train_step --- fl4health/clients/basic_client.py | 2 +- fl4health/clients/nnunet_client.py | 2 +- .../mixins/adaptive_drift_constrained.py | 98 ++++++++++++------- fl4health/mixins/core_protocols.py | 6 ++ 4 files changed, 70 insertions(+), 38 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 1d0f01e06..dd98aad9e 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -579,7 +579,7 @@ def _train_step( # Compute backward pass and update parameters with optimizer losses.backward["backward"].backward() self.transform_gradients(losses) - self.optimizers["global"].step() + optimizer.step() return losses, preds diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 322f2d8c0..73fdaf85d 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -211,7 +211,7 @@ def _train_step( ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": - return super().train_step(input, target) + return super()._train_step(model, optimizer, input, target) # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler # Clear gradients from optimizer if they exist diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py index 841bd0d7e..fbc82133e 100644 --- a/fl4health/mixins/adaptive_drift_constrained.py +++ b/fl4health/mixins/adaptive_drift_constrained.py @@ -1,7 +1,7 @@ """AdaptiveDriftConstrainedMixin""" import warnings -from logging import INFO, WARN, DEBUG +from logging import DEBUG, INFO, WARN from typing import Any, Protocol, runtime_checkable import torch @@ -16,7 +16,7 @@ from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint from fl4health.utils.losses import TrainingLosses -from fl4health.utils.typing import TorchFeatureType, TorchPredType, TorchTargetType +from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType @runtime_checkable @@ -29,6 +29,13 @@ class AdaptiveDriftConstrainedProtocol(BasicClientProtocol, Protocol): def compute_penalty_loss(self) -> torch.Tensor: ... # noqa: E704 + def __compute_training_loss( # noqa: E704 + self, + preds: TorchPredType, + features: TorchFeatureType, + target: TorchTargetType, + ) -> TrainingLosses: ... + class AdaptiveDriftConstrainedMixin: def __init__(self, *args: Any, **kwargs: Any): @@ -148,40 +155,59 @@ def set_parameters( super().set_parameters(server_model_state, config, fitting_round) # type: ignore[safe-super] - # def compute_training_loss( - # self: AdaptiveDriftConstrainedProtocol, - # preds: TorchPredType, - # features: TorchFeatureType, - # target: TorchTargetType, - # ) -> TrainingLosses: - # """ - # Computes training loss given predictions of the model and ground truth data. Adds to objective by including - # penalty loss. - - # Args: - # preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in - # dictionary will be used to compute metrics. - # features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. - # target: (TorchTargetType): Ground truth data to evaluate predictions against. - - # Returns: - # TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed - # by name. Additional losses includes penalty loss. - # """ - # log(DEBUG, "COMPUTE TRAINING LOSS CALL FROM ADAPTIVE MIXIN") - # loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) - # if additional_losses is None: - # additional_losses = {} - - # additional_losses["loss"] = loss.clone() - # # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation - # additional_losses["loss_for_adaptation"] = loss.clone() - - # # Compute the drift penalty loss and store it in the additional losses dictionary. - # penalty_loss = self.compute_penalty_loss() - # additional_losses["penalty_loss"] = penalty_loss.clone() - - # return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) + def train_step( + self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + + # Clear gradients from optimizer if they exist + self.optimizers["global"].zero_grad() + + # Call user defined methods to get predictions and compute loss + preds, features = self.predict(input) + target = self.transform_target(target) + losses = self.__compute_training_loss(preds, features, target) + + # Compute backward pass and update parameters with optimizer + losses.backward["backward"].backward() + self.transform_gradients(losses) + self.optimizers["global"].step() + + return losses, preds + + def __compute_training_loss( + self: AdaptiveDriftConstrainedProtocol, + preds: TorchPredType, + features: TorchFeatureType, + target: TorchTargetType, + ) -> TrainingLosses: + """ + Computes training loss given predictions of the model and ground truth data. Adds to objective by including + penalty loss. + + Args: + preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in + dictionary will be used to compute metrics. + features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. + target: (TorchTargetType): Ground truth data to evaluate predictions against. + + Returns: + TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed + by name. Additional losses includes penalty loss. + """ + log(DEBUG, "COMPUTE TRAINING LOSS CALL FROM ADAPTIVE MIXIN") + loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) + if additional_losses is None: + additional_losses = {} + + additional_losses["loss"] = loss.clone() + # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation + additional_losses["loss_for_adaptation"] = loss.clone() + + # Compute the drift penalty loss and store it in the additional losses dictionary. + penalty_loss = self.compute_penalty_loss() + additional_losses["penalty_loss"] = penalty_loss.clone() + + return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) def get_parameter_exchanger(self: AdaptiveDriftConstrainedProtocol, config: Config) -> ParameterExchanger: """ diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index df7c62a29..f0ab89a1e 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -82,12 +82,18 @@ def _train_step( ) -> tuple[TrainingLosses, TorchPredType]: pass # pragma: no cover + def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + pass # pragma: no cover + def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: pass # pragma: no cover def transform_target(self, target: TorchTargetType) -> TorchTargetType: pass # pragma: no cover + def transform_gradients(self, losses: TrainingLosses) -> None: + pass # pragma: no cover + def compute_training_loss( self, preds: TorchPredType, From 032ac12466b76801e07c8c2c248033bf2350ba1c Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Tue, 27 May 2025 11:32:41 -0400 Subject: [PATCH 09/32] working --- fl4health/clients/basic_client.py | 29 ++++++++++---- fl4health/clients/nnunet_client.py | 20 +++++++++- fl4health/mixins/personalized/ditto.py | 53 +++++++++++++------------- 3 files changed, 66 insertions(+), 36 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index dd98aad9e..10432c930 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -559,15 +559,9 @@ def update_metric_manager( """ metric_manager.update(preds, target) - def _train_step( - self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + def _train_step_compute_preds_and_losses( + self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: - """Helper train step. - - This interface allows for injection of model and optimizer params, which - are useful for personalized FL methods. - """ - # Clear gradients from optimizer if they exist optimizer.zero_grad() @@ -576,11 +570,30 @@ def _train_step( target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) + return losses, preds + + def _train_step_apply_backwards_and_step( + self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses + ) -> tuple[TrainingLosses, TorchPredType]: # Compute backward pass and update parameters with optimizer losses.backward["backward"].backward() self.transform_gradients(losses) optimizer.step() + return losses + + def _train_step( + self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + """Helper train step. + + This interface allows for injection of model and optimizer params, which + are useful for personalized FL methods. + """ + + losses, preds = self._train_step_compute_preds_and_losses(model, optimizer, input, target) + losses = self._train_step_apply_backwards_and_step(model, optimizer, losses) + return losses, preds def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 73fdaf85d..e7bf65025 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -206,7 +206,7 @@ def __init__( if self.verbose: log(INFO, "Disabling model optimizations and JIT compilation. This may impact runtime performance.") - def _train_step( + def _train_step_compute_preds_and_losses( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. @@ -222,6 +222,11 @@ def _train_step( target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) + return losses, preds + + def _train_step_apply_backwards_and_step( + self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses + ) -> tuple[TrainingLosses, TorchPredType]: # Compute scaled loss and perform backward pass scaled_backward_loss = self.grad_scaler.scale(losses.backward["backward"]) scaled_backward_loss.backward() @@ -234,6 +239,19 @@ def _train_step( self.grad_scaler.step(optimizer) self.grad_scaler.update() + return losses + + def _train_step( + self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + """Helper train step. + + This interface allows for injection of model and optimizer params, which + are useful for personalized FL methods. + """ + losses, preds = self._train_step_compute_preds_and_losses(model, optimizer, input, target) + losses = self._train_step_apply_backwards_and_step(model, optimizer, losses) + return losses, preds def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index 94c25d025..a5ce757dc 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -373,49 +373,48 @@ def train_step( """ # global - global_losses, global_preds = self._train_step( + global_losses, global_preds = self._train_step_compute_preds_and_losses( self.safe_global_model(), self.optimizers["global"], input, target ) # local - local_losses, local_preds = self._train_step(self.model, self.optimizers["local"], input, target) + local_losses, local_preds = self._train_step_compute_preds_and_losses( + self.model, self.optimizers["local"], input, target + ) log(DEBUG, f"global_losses: {type(global_losses)}") log(DEBUG, f"local_losses: {type(local_losses)}") - local_losses = cast(TrainingLosses, local_losses) global_losses = cast(TrainingLosses, global_losses) + local_losses = cast(TrainingLosses, local_losses) + local_loss_clone = local_losses.backward["backward"].clone() - # aggregate global and local losses - - local_loss = local_losses.backward["backward"] - global_loss = global_losses.backward["backward"] + # take step + # global + global_losses = self._train_step_apply_backwards_and_step( + self.safe_global_model(), self.optimizers["global"], global_losses + ) + # local + penalty_loss = self.compute_penalty_loss() + local_losses.backward["backward"] = local_losses.backward["backward"] + penalty_loss + local_losses = self._train_step_apply_backwards_and_step(self.model, self.optimizers["local"], local_losses) - # from compute_loss_and_additional_losses + # prepare return values additional_losses = { - "local_loss": local_loss.clone(), - "global_loss": global_loss, + "penalty_loss": penalty_loss.clone(), + "local_loss": local_loss_clone, + "global_loss": global_losses.backward["backward"], + "loss_for_adaptation": local_loss_clone.clone(), } + local_losses.additional_losses = additional_losses - # from compute_loss - # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or - # less weight is used to constrain it to the global model (as in FedProx) - additional_losses["loss_for_adaptation"] = additional_losses["local_loss"].clone() - - # This is the Ditto penalty loss of the local model compared with the original Global model weights, scaled - # by drift_penalty_weight (or lambda in the original paper) - penalty_loss = self.compute_penalty_loss() - additional_losses["penalty_loss"] = penalty_loss.clone() - - training_losses = TrainingLosses(backward=local_loss + penalty_loss, additional_losses=additional_losses) - - # aggregate preds + # combined preds if isinstance(global_preds, torch.Tensor) and isinstance(local_preds, torch.Tensor): - agg_preds = {"global": global_preds, "local": local_preds} + combined_preds = {"global": global_preds, "local": local_preds} elif isinstance(global_preds, dict) and isinstance(local_preds, dict): - agg_preds = {f"global-{k}": v for k, v in global_preds.items()} - agg_preds.update(**{f"local-{k}": v for k, v in local_preds.items()}) + combined_preds = {f"global-{k}": v for k, v in global_preds.items()} + combined_preds.update(**{f"local-{k}": v for k, v in local_preds.items()}) - return training_losses, agg_preds + return local_losses, combined_preds def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Helper method to extract predictions from global and local models. From c7acb7724137b895d8ed6f30304a625fd7ce85b9 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 10:16:59 -0400 Subject: [PATCH 10/32] simple ditto example working --- fl4health/clients/basic_client.py | 21 +++++++++++++++++++-- fl4health/mixins/personalized/ditto.py | 15 +++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 10432c930..55611047a 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -612,7 +612,9 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr """ return self._train_step(self.model, self.optimizers["global"], input, target) - def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + def _val_step( + self, model: nn.Module, input: TorchInputType, target: TorchTargetType + ) -> tuple[EvaluationLosses, TorchPredType]: """ Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already. @@ -627,12 +629,27 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval # Get preds and compute loss with torch.no_grad(): - preds, features = self.predict(input) + preds, features = self._predict(model, input) target = self.transform_target(target) losses = self.compute_evaluation_loss(preds, features, target) return losses, preds + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + """ + Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already. + + Args: + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the + predictions produced by the model. + """ + + return self._val_step(self.model, input) + def train_by_epochs( self, epochs: int, diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index a5ce757dc..d608da1f3 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -416,6 +416,21 @@ def train_step( return local_losses, combined_preds + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + + # global + global_losses, global_preds = self._val_step(self.safe_global_model(), input, target) + # local + local_losses, local_preds = self._val_step(self.model, input, target) + + # combine + losses = EvaluationLosses( + local_losses.checkpoint, + additional_losses={"global_loss": global_losses.checkpoint, "local_loss": local_losses.checkpoint}, + ) + preds = {"global": global_preds["prediction"], "local": local_preds["prediction"]} + return losses, preds + def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Helper method to extract predictions from global and local models. From 6252414b9e4b0cdca6a53df6ff8406017237fe07 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 12:30:14 -0400 Subject: [PATCH 11/32] add val_step in dittoclient --- fl4health/clients/basic_client.py | 4 +-- fl4health/clients/ditto_client.py | 14 +++++++- fl4health/clients/nnunet_client.py | 4 +-- fl4health/mixins/core_protocols.py | 15 +++++++++ fl4health/mixins/personalized/ditto.py | 45 +++++-------------------- tests/mixins/personalized/test_ditto.py | 22 ------------ 6 files changed, 40 insertions(+), 64 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 55611047a..bbdf555f7 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -574,7 +574,7 @@ def _train_step_compute_preds_and_losses( def _train_step_apply_backwards_and_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses - ) -> tuple[TrainingLosses, TorchPredType]: + ) -> TrainingLosses: # Compute backward pass and update parameters with optimizer losses.backward["backward"].backward() self.transform_gradients(losses) @@ -648,7 +648,7 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval predictions produced by the model. """ - return self._val_step(self.model, input) + return self._val_step(self.model, input, target) def train_by_epochs( self, diff --git a/fl4health/clients/ditto_client.py b/fl4health/clients/ditto_client.py index 03749d4aa..00e4f4a5b 100644 --- a/fl4health/clients/ditto_client.py +++ b/fl4health/clients/ditto_client.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from logging import INFO +from logging import DEBUG, INFO from pathlib import Path import torch @@ -254,6 +254,7 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr # Forward pass on both the global and local models preds, features = self.predict(input) + log(DEBUG, f"PRED KEYS from train_step: {preds.keys()}") target = self.transform_target(target) # Apply transformation (Defaults to identity) # Compute all relevant losses @@ -270,6 +271,15 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr # Return dictionary of predictions where key is used to name respective MetricMeters return losses, preds + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + # Get preds and compute loss + with torch.no_grad(): + preds, features = self.predict(input) + target = self.transform_target(target) + losses = self.compute_evaluation_loss(preds, features, target) + + return losses, preds + def predict( self, input: TorchInputType, @@ -328,6 +338,7 @@ def compute_loss_and_additional_losses( """ # Compute global model vanilla loss + log(DEBUG, f"PRED KEYS from compute_loss_and_additional_losses: {preds.keys()}") assert "global" in preds global_loss = self.criterion(preds["global"], target) @@ -366,6 +377,7 @@ def compute_training_loss( assert self.global_model.training and self.model.training # local loss is stored in loss, global model loss is stored in additional losses. + log(DEBUG, f"PRED KEYS: {preds.keys()}") loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index e7bf65025..608046d9b 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -211,7 +211,7 @@ def _train_step_compute_preds_and_losses( ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": - return super()._train_step(model, optimizer, input, target) + return super()._train_step_compute_preds_and_losses(model, optimizer, input, target) # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler # Clear gradients from optimizer if they exist @@ -226,7 +226,7 @@ def _train_step_compute_preds_and_losses( def _train_step_apply_backwards_and_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses - ) -> tuple[TrainingLosses, TorchPredType]: + ) -> TrainingLosses: # Compute scaled loss and perform backward pass scaled_backward_loss = self.grad_scaler.scale(losses.backward["backward"]) scaled_backward_loss.backward() diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index f0ab89a1e..1fb3ac58c 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -77,11 +77,26 @@ def initialize_all_model_weights(self, parameters: NDArrays, config: Config) -> def update_before_train(self, current_server_round: int) -> None: pass # pragma: no cover + def _train_step_compute_preds_and_losses( + self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType + ) -> tuple[TrainingLosses, TorchPredType]: + pass # pragma: no cover + + def _train_step_apply_backwards_and_step( + self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses + ) -> TrainingLosses: + pass # pragma: no cover + def _train_step( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: pass # pragma: no cover + def _val_step( + self, model: nn.Module, input: TorchInputType, target: TorchTargetType + ) -> tuple[EvaluationLosses, TorchPredType]: + pass # pragma: no cover + def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: pass # pragma: no cover diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index d608da1f3..d6b03fbe9 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -1,7 +1,7 @@ """Ditto Personalized Mixin""" import warnings -from logging import DEBUG, INFO, WARN +from logging import INFO, WARN from typing import Any, Protocol, cast, runtime_checkable import torch @@ -34,9 +34,6 @@ def _copy_optimizer_with_new_params(self, original_optimizer: Optimizer) -> Opti def set_initial_global_tensors(self) -> None: pass # pragma: no cover - def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - pass # pragma: no cover - def safe_global_model(self) -> nn.Module: pass # pragma: no cover @@ -74,8 +71,6 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: if not isinstance(self, BasicClientProtocolPreSetup): raise RuntimeError("This object needs to satisfy `BasicClientProtocolPreSetup`.") - # - def __init_subclass__(cls, **kwargs: Any) -> None: """This method is called when a class inherits from AdaptiveMixin""" super().__init_subclass__(**kwargs) @@ -380,13 +375,7 @@ def train_step( local_losses, local_preds = self._train_step_compute_preds_and_losses( self.model, self.optimizers["local"], input, target ) - - log(DEBUG, f"global_losses: {type(global_losses)}") - log(DEBUG, f"local_losses: {type(local_losses)}") - - global_losses = cast(TrainingLosses, global_losses) - local_losses = cast(TrainingLosses, local_losses) - local_loss_clone = local_losses.backward["backward"].clone() + local_loss_clone = local_losses.backward["backward"].clone() # need a clone for later # take step # global @@ -416,7 +405,9 @@ def train_step( return local_losses, combined_preds - def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + def val_step( + self: DittoPersonalizedProtocol, input: TorchInputType, target: TorchTargetType + ) -> tuple[EvaluationLosses, TorchPredType]: # global global_losses, global_preds = self._val_step(self.safe_global_model(), input, target) @@ -428,31 +419,11 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval local_losses.checkpoint, additional_losses={"global_loss": global_losses.checkpoint, "local_loss": local_losses.checkpoint}, ) - preds = {"global": global_preds["prediction"], "local": local_preds["prediction"]} + preds: TorchPredType = {} + preds.update(**{f"global-{k}": v for k, v in global_preds.items()}) + preds.update(**{f"local-{k}": v for k, v in local_preds.items()}) return losses, preds - def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Helper method to extract predictions from global and local models. - - Args: - kind (str): the kind of predictions to be extracted i.e., for 'local' or for 'global' - preds (dict[str, torch.Tensor]): the full preds result from `self.predict()` - - Raises: - ValueError: If supplied `kind` is not of 'global' or 'local' - - Returns: - dict[str, torch.Tensor]: the predictions associated with the specified kind - """ - 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 - @ensure_protocol_compliance def validate( self: DittoPersonalizedProtocol, include_losses_in_metrics: bool = False diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py index e4b464ca9..ed1bc7a71 100644 --- a/tests/mixins/personalized/test_ditto.py +++ b/tests/mixins/personalized/test_ditto.py @@ -303,28 +303,6 @@ def test_predict_delagation(private_predict: MagicMock) -> None: assert_close(res["local"], torch.ones(5)) -def test_extract_pred() -> None: - # setup client - client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) - - res = client._extract_pred( - kind="global", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)} - ) - - assert_close(res["xyz"], torch.ones(5)) - assert_close(res["abc"], torch.zeros(5)) - - -def test_extract_pred_raises_error() -> None: - # setup client - client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) - - with pytest.raises(ValueError): - client._extract_pred( - kind="oops", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)} - ) - - @patch.object(_TestDittoedClient, "set_initial_global_tensors") @patch.object(_TestBasicClient, "update_before_train") def test_update_before_train( From 938e976d21e71a0095507c211908c0794437b542 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 13:05:15 -0400 Subject: [PATCH 12/32] update and simplify adaptive_drift_constrained mixin --- .../mixins/adaptive_drift_constrained.py | 69 ++++--------------- 1 file changed, 15 insertions(+), 54 deletions(-) diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py index fbc82133e..b1e3038f4 100644 --- a/fl4health/mixins/adaptive_drift_constrained.py +++ b/fl4health/mixins/adaptive_drift_constrained.py @@ -1,7 +1,7 @@ """AdaptiveDriftConstrainedMixin""" import warnings -from logging import DEBUG, INFO, WARN +from logging import INFO, WARN from typing import Any, Protocol, runtime_checkable import torch @@ -16,7 +16,7 @@ from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint from fl4health.utils.losses import TrainingLosses -from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType +from fl4health.utils.typing import TorchInputType, TorchPredType, TorchTargetType @runtime_checkable @@ -29,13 +29,6 @@ class AdaptiveDriftConstrainedProtocol(BasicClientProtocol, Protocol): def compute_penalty_loss(self) -> torch.Tensor: ... # noqa: E704 - def __compute_training_loss( # noqa: E704 - self, - preds: TorchPredType, - features: TorchFeatureType, - target: TorchTargetType, - ) -> TrainingLosses: ... - class AdaptiveDriftConstrainedMixin: def __init__(self, *args: Any, **kwargs: Any): @@ -159,56 +152,24 @@ def train_step( self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: - # Clear gradients from optimizer if they exist - self.optimizers["global"].zero_grad() + losses, preds = self._train_step_compute_preds_and_losses(self.model, self.optimizers["global"], input, target) + loss_clone = losses.backward["backward"].clone() - # Call user defined methods to get predictions and compute loss - preds, features = self.predict(input) - target = self.transform_target(target) - losses = self.__compute_training_loss(preds, features, target) + # apply penalty + penalty_loss = self.compute_penalty_loss() + losses.backward["backward"] = losses.backward["backward"] + penalty_loss + losses = self._train_step_apply_backwards_and_step(self.model, self.optimizers["global"], losses) - # Compute backward pass and update parameters with optimizer - losses.backward["backward"].backward() - self.transform_gradients(losses) - self.optimizers["global"].step() + # prepare return values + additional_losses = { + "penalty_loss": penalty_loss.clone(), + "local_loss": loss_clone, + "loss_for_adaptation": loss_clone.clone(), + } + losses.additional_losses = additional_losses return losses, preds - def __compute_training_loss( - self: AdaptiveDriftConstrainedProtocol, - preds: TorchPredType, - features: TorchFeatureType, - target: TorchTargetType, - ) -> TrainingLosses: - """ - Computes training loss given predictions of the model and ground truth data. Adds to objective by including - penalty loss. - - Args: - preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in - dictionary will be used to compute metrics. - features: (TorchFeatureType): Feature(s) of the model(s) indexed by name. - target: (TorchTargetType): Ground truth data to evaluate predictions against. - - Returns: - TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed - by name. Additional losses includes penalty loss. - """ - log(DEBUG, "COMPUTE TRAINING LOSS CALL FROM ADAPTIVE MIXIN") - loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target) - if additional_losses is None: - additional_losses = {} - - additional_losses["loss"] = loss.clone() - # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation - additional_losses["loss_for_adaptation"] = loss.clone() - - # Compute the drift penalty loss and store it in the additional losses dictionary. - penalty_loss = self.compute_penalty_loss() - additional_losses["penalty_loss"] = penalty_loss.clone() - - return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses) - def get_parameter_exchanger(self: AdaptiveDriftConstrainedProtocol, config: Config) -> ParameterExchanger: """ Setting up the parameter exchanger to include the appropriate packing functionality. From a48ededc3b18c2d4032b6f20644c2144ca2879b7 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 13:16:17 -0400 Subject: [PATCH 13/32] rm old tests that are outdated --- tests/mixins/personalized/test_ditto.py | 112 +----------------------- 1 file changed, 4 insertions(+), 108 deletions(-) diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py index ed1bc7a71..9bb0f6e16 100644 --- a/tests/mixins/personalized/test_ditto.py +++ b/tests/mixins/personalized/test_ditto.py @@ -11,7 +11,8 @@ from numpy.testing import assert_array_equal from torch.nn.modules.loss import _Loss from torch.optim import Optimizer -from torch.testing import assert_close + +# from torch.testing import assert_close from torch.utils.data import DataLoader, TensorDataset from fl4health.clients.basic_client import BasicClient @@ -27,7 +28,8 @@ from fl4health.parameter_exchange.parameter_packer import ( ParameterPackerAdaptiveConstraint, ) -from fl4health.utils.losses import TrainingLosses + +# from fl4health.utils.losses import TrainingLosses from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType @@ -239,70 +241,6 @@ def test_get_optimizer(mock_copy_optimizer: MagicMock) -> None: mock_copy_optimizer.assert_called_once_with(client.optimizers["local"]) -def test_predict() -> None: - # setup client - client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) - - mock_model = MagicMock() - mock_global_model = MagicMock() - - mock_model.return_value = torch.ones(5) - mock_global_model.return_value = torch.zeros(5) - - client.model = mock_model - client.global_model = mock_global_model - - client.optimizers = { - "global": MagicMock(), - "local": MagicMock(), - } - - client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint()) - client.initialized = True - - # act - # TODO: fix the mixin/protocol typing that leads to mypy complaint - res, _ = client.predict(input=torch.zeros(5)) # type: ignore - print(f"res: {res}") - print(torch.zeros(5)) - - # assert - assert_close(res["global"], torch.zeros(5)) - assert_close(res["local"], torch.ones(5)) - - -@patch.object(_TestDittoedClientV2, "_predict") -def test_predict_delagation(private_predict: MagicMock) -> None: - # setup client - client = _TestDittoedClientV2(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) - client.model = torch.nn.Linear(5, 5) - client.global_model = torch.nn.Linear(5, 5) - - private_predict.side_effect = [(torch.zeros(5), {}), (torch.ones(5), {})] - - client.optimizers = { - "global": MagicMock(), - "local": MagicMock(), - } - - client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint()) - client.initialized = True - - # act - # TODO: fix the mixin/protocol typing that leads to mypy complaint - res, _ = client.predict(input=torch.zeros(5)) # type: ignore - print(f"res: {res}") - print(torch.zeros(5)) - - # assert - assert_close(res["global"], torch.zeros(5)) - assert_close(res["local"], torch.ones(5)) - - @patch.object(_TestDittoedClient, "set_initial_global_tensors") @patch.object(_TestBasicClient, "update_before_train") def test_update_before_train( @@ -326,45 +264,3 @@ def test_safe_model_raises_error() -> None: with pytest.raises(ValueError): # TODO: fix the mixin/protocol typing that leads to mypy complaint client.safe_global_model() # type: ignore - - -@patch("torch.optim.Optimizer") -@patch.object(_TestDittoedClient, "predict") -@patch.object(_TestDittoedClient, "compute_training_loss") -def test_train_step( - mock_compute_training_loss: MagicMock, mock_predict: MagicMock, mock_optimizer_class: MagicMock -) -> None: - # setup client - client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) - client.optimizers = { - "global": MagicMock(), - "local": MagicMock(), - } - client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000)))) - client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint()) - client.initialized = True - - # arrange mocks - pred_return: tuple[TorchPredType, TorchFeatureType] = {"local": torch.ones(5)}, {} - preds, features = pred_return - mock_predict.return_value = preds, features - mock_backward_loss = MagicMock() - mock_additional_global_loss = MagicMock() - losses = TrainingLosses( - backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss} - ) - mock_compute_training_loss.return_value = losses - - # act - input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3) - # TODO: fix the mixin/protocol typing that leads to mypy complaint - retval = client.train_step(input, target) # type: ignore - - mock_predict.assert_called_once() - client.optimizers["global"].zero_grad.assert_called_once() - client.optimizers["local"].zero_grad.assert_called_once() - mock_compute_training_loss.assert_called_once_with(preds, features, target) - mock_backward_loss.backward.assert_called_once() - mock_additional_global_loss.backward.assert_called_once() - assert retval == (losses, preds) From 007c1794abce582cf078bed86960082dc962bea4 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 13:22:55 -0400 Subject: [PATCH 14/32] shawn cr for get_global_model --- fl4health/clients/nnunet_client.py | 10 ++-------- fl4health/mixins/personalized/ditto.py | 6 +++--- fl4health/utils/config.py | 2 -- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 608046d9b..5eb660632 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -1,4 +1,3 @@ -import copy import gc import logging import os @@ -28,7 +27,7 @@ from fl4health.metrics.base_metrics import Metric from fl4health.metrics.metric_managers import MetricManager from fl4health.reporting.base_reporter import BaseReporter -from fl4health.utils.config import FOR_GLOBAL_MODEL_KEY, narrow_dict_type +from fl4health.utils.config import narrow_dict_type from fl4health.utils.logging import LoggingMode from fl4health.utils.losses import LossMeterType, TrainingLosses from fl4health.utils.nnunet_utils import ( @@ -340,12 +339,7 @@ def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]: return train_loader, val_loader def get_model(self, config: Config) -> nn.Module: - - for_global = config.get(FOR_GLOBAL_MODEL_KEY, False) - if for_global: - return copy.deepcopy(self.nnunet_trainer.network) - else: - return self.nnunet_trainer.network + return self.nnunet_trainer.network def get_criterion(self, config: Config) -> _Loss: if isinstance(self.nnunet_trainer.loss, DeepSupervisionWrapper): diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index d6b03fbe9..d8c1beaf0 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -1,5 +1,6 @@ """Ditto Personalized Mixin""" +import copy import warnings from logging import INFO, WARN from typing import Any, Protocol, cast, runtime_checkable @@ -15,7 +16,7 @@ from fl4health.mixins.core_protocols import BasicClientProtocolPreSetup from fl4health.mixins.personalized.utils import ensure_protocol_compliance from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger -from fl4health.utils.config import FOR_GLOBAL_MODEL_KEY, narrow_dict_type +from fl4health.utils.config import narrow_dict_type from fl4health.utils.losses import EvaluationLosses, TrainingLosses from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType, TorchTargetType @@ -172,8 +173,7 @@ def get_global_model(self: DittoPersonalizedProtocol, config: Config) -> nn.Modu Returns: nn.Module: The PyTorch model serving as the global model for Ditto """ - config[FOR_GLOBAL_MODEL_KEY] = True - return self.get_model(config).to(self.device) + return copy.deepcopy(self.get_model(config).to(self.device)) @ensure_protocol_compliance def get_optimizer(self: DittoPersonalizedProtocol, config: Config) -> dict[str, Optimizer]: diff --git a/fl4health/utils/config.py b/fl4health/utils/config.py index 8beb2b70d..7a8becaf9 100644 --- a/fl4health/utils/config.py +++ b/fl4health/utils/config.py @@ -8,8 +8,6 @@ "batch_size": int, } -FOR_GLOBAL_MODEL_KEY = "__for_global_model" - T = TypeVar("T") From d3f84c81c6c018669d1fec8aa9ef27e1b2ad55f4 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Wed, 28 May 2025 13:28:07 -0400 Subject: [PATCH 15/32] shawns cr on notes about update_metric_manager --- fl4health/clients/nnunet_client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 5eb660632..bc29069e0 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -780,10 +780,11 @@ def update_metric_manager( log(DEBUG, f"preds: {preds.keys()}") - # If personalized, the + # If personalized preds will have keys prefixed with "global-" and "local-", + # but we only want to keep the local ones here. NOTE: this will still work + # even in the non-personalized case as those preds don't get prefixed with "global-" preds = {k: v for k, v in preds.items() if not k.startswith("global")} - # remove prefix - preds = {k.replace("local-", ""): v for k, v in preds.items()} + preds = {k.replace("local-", ""): v for k, v in preds.items()} # no-op in non-personalized case if len(preds) > 1: # for nnunet the first pred in the output list is the main one From 6d88019c790f6e0b140be98bd047d593f8793e53 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Thu, 29 May 2025 01:06:49 -0400 Subject: [PATCH 16/32] perfcl add val step since we override predict --- fl4health/clients/perfcl_client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/fl4health/clients/perfcl_client.py b/fl4health/clients/perfcl_client.py index 251083bd0..0eae1b282 100644 --- a/fl4health/clients/perfcl_client.py +++ b/fl4health/clients/perfcl_client.py @@ -167,6 +167,15 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp return preds, features + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: + # Get preds and compute loss + with torch.no_grad(): + preds, features = self.predict(input) + target = self.transform_target(target) + losses = self.compute_evaluation_loss(preds, features, target) + + return losses, preds + def update_after_train(self, local_steps: int, loss_dict: dict[str, float], config: Config) -> None: """ This function is called after client-side training concludes. In this case, it is used to save the local From 5d3a05edee92f3187db5f6af6eec01febeb4782d Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Thu, 29 May 2025 12:05:06 -0400 Subject: [PATCH 17/32] perfcl client _predict --- fl4health/clients/perfcl_client.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/fl4health/clients/perfcl_client.py b/fl4health/clients/perfcl_client.py index 0eae1b282..69167ef12 100644 --- a/fl4health/clients/perfcl_client.py +++ b/fl4health/clients/perfcl_client.py @@ -2,6 +2,7 @@ from pathlib import Path import torch +import torch.nn as nn from flwr.common.typing import Config from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule @@ -136,7 +137,7 @@ def _all_contrastive_loss_modules_defined(self) -> bool: and self.initial_global_module is not None ) - def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: """ Computes the prediction(s) and features of the model(s) given the input. @@ -167,6 +168,22 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp return preds, features + def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + """ + Computes the prediction(s) and features of the model(s) given the input. + + Args: + input (TorchInputType): Inputs to be fed into the model. ``TorchInputType`` is simply an alias + for the union of ``torch.Tensor`` and ``dict[str, torch.Tensor]``. + + Returns: + tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: A tuple in which the first element + contains predictions indexed by name and the second element contains intermediate activations + index by name. Specifically the features of the model, features of the global model and features of + the old model are returned. All predictions included in dictionary will be used to compute metrics. + """ + return self._predict(self.model, input) + def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: # Get preds and compute loss with torch.no_grad(): From 09a08f930052d813647635e0afca8e4d5baa59d9 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 01:25:34 -0400 Subject: [PATCH 18/32] cleanup --- fl4health/clients/nnunet_client.py | 13 ------------- fl4health/clients/perfcl_client.py | 11 +---------- 2 files changed, 1 insertion(+), 23 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index bc29069e0..e873d7aa0 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -240,19 +240,6 @@ def _train_step_apply_backwards_and_step( return losses - def _train_step( - self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType - ) -> tuple[TrainingLosses, TorchPredType]: - """Helper train step. - - This interface allows for injection of model and optimizer params, which - are useful for personalized FL methods. - """ - losses, preds = self._train_step_compute_preds_and_losses(model, optimizer, input, target) - losses = self._train_step_apply_backwards_and_step(model, optimizer, losses) - - return losses, preds - def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: """ Given a single batch of input and target data, generate predictions, compute loss, update parameters and diff --git a/fl4health/clients/perfcl_client.py b/fl4health/clients/perfcl_client.py index 69167ef12..0ccd49173 100644 --- a/fl4health/clients/perfcl_client.py +++ b/fl4health/clients/perfcl_client.py @@ -153,7 +153,7 @@ def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredTy """ # For PerFCL models, we required the input to simply be a torch.Tensor assert isinstance(input, torch.Tensor) - preds, features = self.model(input) + preds, features = model(input) # In the first server round, these module will not have been set. if ( self.old_local_module is not None @@ -184,15 +184,6 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp """ return self._predict(self.model, input) - def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]: - # Get preds and compute loss - with torch.no_grad(): - preds, features = self.predict(input) - target = self.transform_target(target) - losses = self.compute_evaluation_loss(preds, features, target) - - return losses, preds - def update_after_train(self, local_steps: int, loss_dict: dict[str, float], config: Config) -> None: """ This function is called after client-side training concludes. In this case, it is used to save the local From 7bdb974a46decf7eae97203a3f03f86308e4210f Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 01:38:29 -0400 Subject: [PATCH 19/32] _predict rename to _predict_with_model --- fl4health/clients/basic_client.py | 30 ++++++++++++---- fl4health/clients/nnunet_client.py | 46 +++++++++++-------------- fl4health/clients/perfcl_client.py | 4 +-- fl4health/mixins/core_protocols.py | 2 +- tests/mixins/personalized/test_ditto.py | 4 ++- 5 files changed, 50 insertions(+), 36 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index bbdf555f7..263daddb6 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -566,7 +566,7 @@ def _train_step_compute_preds_and_losses( optimizer.zero_grad() # Call user defined methods to get predictions and compute loss - preds, features = self._predict(model, input) + preds, features = self._predict_with_model(model, input) target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) @@ -629,7 +629,7 @@ def _val_step( # Get preds and compute loss with torch.no_grad(): - preds, features = self._predict(model, input) + preds, features = self._predict_with_model(model, input) target = self.transform_target(target) losses = self.compute_evaluation_loss(preds, features, target) @@ -1016,10 +1016,28 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger: """ return FullParameterExchanger() - def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: - """Helper predict method. + def _predict_with_model( + self, model: torch.nn.Module, input: TorchInputType + ) -> tuple[TorchPredType, TorchFeatureType]: + """Helper predict interface allowing for injection of model. - Unlike, predict(), this interface allows for injecting the model param. + Unlike, predict(), this interface allows for a model to be supplied. + Subclasses should implement this method if there is need to specialize + the predict method of the client. + + Args: + model (torch.nn.Module): the model with which to make predictions + input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``, + it is assumed that the keys of input match the names of the keyword arguments of + ``self.model.forward().` + + Res: + TypeError: _description_ + ValueError: _description_ + ValueError: _description_ + + Returns: + tuple[TorchPredType, TorchFeatureType]: _description_ """ if isinstance(input, torch.Tensor): @@ -1065,7 +1083,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model forward. """ - return self._predict(self.model, input) + return self._predict_with_model(self.model, input) def compute_loss_and_additional_losses( self, preds: TorchPredType, features: TorchFeatureType, target: TorchTargetType diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index e873d7aa0..3bf0c3c8c 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -217,7 +217,7 @@ def _train_step_compute_preds_and_losses( optimizer.zero_grad() # Call user defined methods to get predictions and compute loss - preds, features = self._predict(model, input) + preds, features = self._predict_with_model(model, input) target = self.transform_target(target) losses = self.compute_training_loss(preds, features, target) @@ -617,7 +617,9 @@ def setup_client(self, config: Config) -> None: # We have to call parent method after setting up nnunet trainer super().setup_client(config) - def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]: + def _predict_with_model( + self, model: torch.nn.Module, input: TorchInputType + ) -> tuple[TorchPredType, dict[str, torch.Tensor]]: if isinstance(input, torch.Tensor): # If device type is cuda, nnUNet defaults to mixed precision forward pass if self.device.type == "cuda": @@ -654,14 +656,28 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch tuple[TorchPredType, dict[str, torch.Tensor]]: A tuple in which the first element model outputs indexed by name. The second element is unused by this subclass and therefore is always an empty dict """ - return self._predict(self.model, input) + return self._predict_with_model(self.model, input) - def _special_compute_loss_and_additional_losses( + def compute_loss_and_additional_losses( self, preds: TorchPredType, features: dict[str, torch.Tensor], target: TorchTargetType, ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + """ + Checks the pred and target types and computes the loss. If device type is cuda, loss computed in mixed + precision. + + Args: + preds (TorchPredType): Dictionary of model output tensors indexed by name + features (dict[str, torch.Tensor]): Not used by this subclass + target (TorchTargetType): The targets to evaluate the predictions with. If multiple prediction tensors + are given, target must be a dictionary with the same number of tensors + + Returns: + tuple[torch.Tensor, dict[str, torch.Tensor] | None]: A tuple where the first element is the loss and the + second element is an optional additional loss + """ # If deep supervision is turned on we must convert loss and target dicts into lists loss_preds = prepare_loss_arg(preds) loss_targets = prepare_loss_arg(target) @@ -686,28 +702,6 @@ def _special_compute_loss_and_additional_losses( return loss - def compute_loss_and_additional_losses( - self, - preds: TorchPredType, - features: dict[str, torch.Tensor], - target: TorchTargetType, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: - """ - Checks the pred and target types and computes the loss. If device type is cuda, loss computed in mixed - precision. - - Args: - preds (TorchPredType): Dictionary of model output tensors indexed by name - features (dict[str, torch.Tensor]): Not used by this subclass - target (TorchTargetType): The targets to evaluate the predictions with. If multiple prediction tensors - are given, target must be a dictionary with the same number of tensors - - Returns: - tuple[torch.Tensor, dict[str, torch.Tensor] | None]: A tuple where the first element is the loss and the - second element is an optional additional loss - """ - return self._special_compute_loss_and_additional_losses(preds, features, target) - def mask_data(self, pred: torch.Tensor, target: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """ Masks the pred and target tensors according to nnunet ``ignore_label``. The number of classes in the input diff --git a/fl4health/clients/perfcl_client.py b/fl4health/clients/perfcl_client.py index 0ccd49173..0b50337ee 100644 --- a/fl4health/clients/perfcl_client.py +++ b/fl4health/clients/perfcl_client.py @@ -137,7 +137,7 @@ def _all_contrastive_loss_modules_defined(self) -> bool: and self.initial_global_module is not None ) - def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + def _predict_with_model(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: """ Computes the prediction(s) and features of the model(s) given the input. @@ -182,7 +182,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp index by name. Specifically the features of the model, features of the global model and features of the old model are returned. All predictions included in dictionary will be used to compute metrics. """ - return self._predict(self.model, input) + return self._predict_with_model(self.model, input) def update_after_train(self, local_steps: int, loss_dict: dict[str, float], config: Config) -> None: """ diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index 1fb3ac58c..21042f35d 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -97,7 +97,7 @@ def _val_step( ) -> tuple[EvaluationLosses, TorchPredType]: pass # pragma: no cover - def _predict(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + def _predict_with_model(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: pass # pragma: no cover def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py index 9bb0f6e16..76a04e41b 100644 --- a/tests/mixins/personalized/test_ditto.py +++ b/tests/mixins/personalized/test_ditto.py @@ -60,7 +60,9 @@ def get_optimizer(self, config: Config) -> Optimizer | dict[str, Optimizer]: def get_criterion(self, config: Config) -> _Loss: return torch.nn.CrossEntropyLoss() - def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]: + def _predict_with_model( + self, model: torch.nn.Module, input: TorchInputType + ) -> tuple[TorchPredType, TorchFeatureType]: return {}, {} From be543c065ecb5b44468c712a679ea0364509cc8b Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 01:44:55 -0400 Subject: [PATCH 20/32] rename _train_step --- fl4health/clients/basic_client.py | 18 +++++++++++++----- fl4health/clients/nnunet_client.py | 2 +- fl4health/mixins/core_protocols.py | 2 +- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 263daddb6..24bcb4f59 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -582,13 +582,21 @@ def _train_step_apply_backwards_and_step( return losses - def _train_step( + def _train_step_with_model_and_optimizer( self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: - """Helper train step. + """Helper train step method that allows for injection of model and optimizer. - This interface allows for injection of model and optimizer params, which - are useful for personalized FL methods. + Subclasses should implement this method if there is a need to specialize + the train_step logic. + + Args: + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with + a dictionary of any predictions produced by the model. """ losses, preds = self._train_step_compute_preds_and_losses(model, optimizer, input, target) @@ -610,7 +618,7 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any predictions produced by the model. """ - return self._train_step(self.model, self.optimizers["global"], input, target) + return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target) def _val_step( self, model: nn.Module, input: TorchInputType, target: TorchTargetType diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 3bf0c3c8c..6932c5429 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -258,7 +258,7 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr predictions produced by the model. """ # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. - return self._train_step(self.model, self.optimizers["global"], input, target) + return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target) @use_default_signal_handlers # Dataloaders use multiprocessing def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]: diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index 21042f35d..b18eb0047 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -87,7 +87,7 @@ def _train_step_apply_backwards_and_step( ) -> TrainingLosses: pass # pragma: no cover - def _train_step( + def _train_step_with_model_and_optimizer( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: pass # pragma: no cover From f95a9e835db5a524f24bdcfd251495fbdc153a47 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:04:19 -0400 Subject: [PATCH 21/32] more private method renames --- fl4health/clients/basic_client.py | 49 +++++++++++++++---- fl4health/clients/nnunet_client.py | 6 +-- .../mixins/adaptive_drift_constrained.py | 4 +- fl4health/mixins/core_protocols.py | 6 +-- fl4health/mixins/personalized/ditto.py | 16 +++--- 5 files changed, 56 insertions(+), 25 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 24bcb4f59..da406b9a1 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -559,9 +559,25 @@ def update_metric_manager( """ metric_manager.update(preds, target) - def _train_step_compute_preds_and_losses( + def _compute_preds_and_losses( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: + """Helper method within the train step for computing preds and losses. + + NOTE: Subclasses should implement this helper method if there is a need + to specialize this part of the overall train step. + + Args: + model (nn.Module): the model used to make predictions + optimizer (Optimizer): the associated optimizer + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with + a dictionary of any predictions produced by the model prior to the + application of the backwards phase + """ # Clear gradients from optimizer if they exist optimizer.zero_grad() @@ -572,9 +588,22 @@ def _train_step_compute_preds_and_losses( return losses, preds - def _train_step_apply_backwards_and_step( + def _apply_backwards_on_losses_and_take_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses ) -> TrainingLosses: + """Helper method within the train step for applying backwards on losses and taking step with optimizer. + + NOTE: Subclasses should implement this helper method if there is a need + to specialize this part of the overall train step. + + Args: + model (nn.Module): the model used for making predictions. Passed here in case subclasses need it. + optimizer (Optimizer): the optimizer with which we take the step + losses (TrainingLosses): the losses to apply backwards on + + Returns: + TrainingLosses: The losses object post backwards application + """ # Compute backward pass and update parameters with optimizer losses.backward["backward"].backward() self.transform_gradients(losses) @@ -587,7 +616,7 @@ def _train_step_with_model_and_optimizer( ) -> tuple[TrainingLosses, TorchPredType]: """Helper train step method that allows for injection of model and optimizer. - Subclasses should implement this method if there is a need to specialize + NOTE: Subclasses should implement this method if there is a need to specialize the train_step logic. Args: @@ -599,8 +628,8 @@ def _train_step_with_model_and_optimizer( a dictionary of any predictions produced by the model. """ - losses, preds = self._train_step_compute_preds_and_losses(model, optimizer, input, target) - losses = self._train_step_apply_backwards_and_step(model, optimizer, losses) + losses, preds = self._compute_preds_and_losses(model, optimizer, input, target) + losses = self._apply_backwards_on_losses_and_take_step(model, optimizer, losses) return losses, preds @@ -620,11 +649,13 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr """ return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target) - def _val_step( + def _val_step_with_model( self, model: nn.Module, input: TorchInputType, target: TorchTargetType ) -> tuple[EvaluationLosses, TorchPredType]: - """ - Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already. + """Helper method for val_step that allows for injection of model. + + NOTE: Subclasses should implement this method if there is a need to + specialize the val_step logic. Args: input (TorchInputType): The input to be fed into the model. @@ -656,7 +687,7 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval predictions produced by the model. """ - return self._val_step(self.model, input, target) + return self._val_step_with_model(self.model, input, target) def train_by_epochs( self, diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 6932c5429..aa492444c 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -205,12 +205,12 @@ def __init__( if self.verbose: log(INFO, "Disabling model optimizations and JIT compilation. This may impact runtime performance.") - def _train_step_compute_preds_and_losses( + def _compute_preds_and_losses( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": - return super()._train_step_compute_preds_and_losses(model, optimizer, input, target) + return super()._compute_preds_and_losses(model, optimizer, input, target) # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler # Clear gradients from optimizer if they exist @@ -223,7 +223,7 @@ def _train_step_compute_preds_and_losses( return losses, preds - def _train_step_apply_backwards_and_step( + def _apply_backwards_on_losses_and_take_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses ) -> TrainingLosses: # Compute scaled loss and perform backward pass diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py index b1e3038f4..ea518f89c 100644 --- a/fl4health/mixins/adaptive_drift_constrained.py +++ b/fl4health/mixins/adaptive_drift_constrained.py @@ -152,13 +152,13 @@ def train_step( self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: - losses, preds = self._train_step_compute_preds_and_losses(self.model, self.optimizers["global"], input, target) + losses, preds = self._compute_preds_and_losses(self.model, self.optimizers["global"], input, target) loss_clone = losses.backward["backward"].clone() # apply penalty penalty_loss = self.compute_penalty_loss() losses.backward["backward"] = losses.backward["backward"] + penalty_loss - losses = self._train_step_apply_backwards_and_step(self.model, self.optimizers["global"], losses) + losses = self._apply_backwards_on_losses_and_take_step(self.model, self.optimizers["global"], losses) # prepare return values additional_losses = { diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index b18eb0047..5d21f23f7 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -77,12 +77,12 @@ def initialize_all_model_weights(self, parameters: NDArrays, config: Config) -> def update_before_train(self, current_server_round: int) -> None: pass # pragma: no cover - def _train_step_compute_preds_and_losses( + def _compute_preds_and_losses( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: pass # pragma: no cover - def _train_step_apply_backwards_and_step( + def _apply_backwards_on_losses_and_take_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses ) -> TrainingLosses: pass # pragma: no cover @@ -92,7 +92,7 @@ def _train_step_with_model_and_optimizer( ) -> tuple[TrainingLosses, TorchPredType]: pass # pragma: no cover - def _val_step( + def _val_step_with_model( self, model: nn.Module, input: TorchInputType, target: TorchTargetType ) -> tuple[EvaluationLosses, TorchPredType]: pass # pragma: no cover diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index d8c1beaf0..42cc0a8c3 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -368,24 +368,24 @@ def train_step( """ # global - global_losses, global_preds = self._train_step_compute_preds_and_losses( + global_losses, global_preds = self._compute_preds_and_losses( self.safe_global_model(), self.optimizers["global"], input, target ) # local - local_losses, local_preds = self._train_step_compute_preds_and_losses( - self.model, self.optimizers["local"], input, target - ) + local_losses, local_preds = self._compute_preds_and_losses(self.model, self.optimizers["local"], input, target) local_loss_clone = local_losses.backward["backward"].clone() # need a clone for later # take step # global - global_losses = self._train_step_apply_backwards_and_step( + global_losses = self._apply_backwards_on_losses_and_take_step( self.safe_global_model(), self.optimizers["global"], global_losses ) # local penalty_loss = self.compute_penalty_loss() local_losses.backward["backward"] = local_losses.backward["backward"] + penalty_loss - local_losses = self._train_step_apply_backwards_and_step(self.model, self.optimizers["local"], local_losses) + local_losses = self._apply_backwards_on_losses_and_take_step( + self.model, self.optimizers["local"], local_losses + ) # prepare return values additional_losses = { @@ -410,9 +410,9 @@ def val_step( ) -> tuple[EvaluationLosses, TorchPredType]: # global - global_losses, global_preds = self._val_step(self.safe_global_model(), input, target) + global_losses, global_preds = self._val_step_with_model(self.safe_global_model(), input, target) # local - local_losses, local_preds = self._val_step(self.model, input, target) + local_losses, local_preds = self._val_step_with_model(self.model, input, target) # combine losses = EvaluationLosses( From 601d95ba3db2532d3889c803e56d5b2ca6118e6d Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:24:49 -0400 Subject: [PATCH 22/32] cr --- examples/nnunet_pfl_example/client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/nnunet_pfl_example/client.py b/examples/nnunet_pfl_example/client.py index 18ee2359a..eab6ebfe3 100644 --- a/examples/nnunet_pfl_example/client.py +++ b/examples/nnunet_pfl_example/client.py @@ -209,7 +209,7 @@ def main( type=str, required=False, default=None, - help="[OPTIONAL] Personalized strategy to use. Can be 'ditto' or 'mr-mtl' \ + help="[OPTIONAL] Personalized strategy to use. For now, can only be 'ditto'. \ Defaults to None, in which no personalized strategy is applied.", ) From 6ba28e02e94c40efe5849b02a0653d3046267c72 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:30:22 -0400 Subject: [PATCH 23/32] nit docstring --- fl4health/clients/basic_client.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index da406b9a1..fdb8f7b1e 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -1058,11 +1058,10 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger: def _predict_with_model( self, model: torch.nn.Module, input: TorchInputType ) -> tuple[TorchPredType, TorchFeatureType]: - """Helper predict interface allowing for injection of model. + """Helper predict method that allows for injection of model. - Unlike, predict(), this interface allows for a model to be supplied. - Subclasses should implement this method if there is need to specialize - the predict method of the client. + NOTE: Subclasses should implement this method if there is need to specialize + the predict logic of the client. Args: model (torch.nn.Module): the model with which to make predictions @@ -1070,13 +1069,17 @@ def _predict_with_model( it is assumed that the keys of input match the names of the keyword arguments of ``self.model.forward().` - Res: - TypeError: _description_ - ValueError: _description_ - ValueError: _description_ - Returns: - tuple[TorchPredType, TorchFeatureType]: _description_ + tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of + predictions indexed by name and the second element contains intermediate activations indexed by name. By + passing features, we can compute losses such as the contrastive loss in MOON. All predictions included in + dictionary will by default be used to compute metrics separately. + + Raises: + TypeError: Occurs when something other than a tensor or dict of tensors is passed in to the model's + forward method. + ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model + forward. """ if isinstance(input, torch.Tensor): From b6a2f36115161ef1179e18bbc2cf1a654ecce147 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:36:02 -0400 Subject: [PATCH 24/32] transform_gradients helper --- fl4health/clients/basic_client.py | 12 ++++++++---- fl4health/clients/nnunet_client.py | 5 +---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index fdb8f7b1e..9cb096e5a 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -1390,12 +1390,16 @@ def update_before_epoch(self, epoch: int) -> None: pass def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: - """ - Helper transform gradients method. + """Helper transform gradients method that allows for injection of model. + + NOTE: Subclasses should implement this helper should there be a need to specialize the logic + for transforming gradients. - Unlike transform_gradients(), this helper's interface allows for injecting - model as a param. + Args: + model (torch.nn.Module): the model used to generate predictions to compute losses + losses (TrainingLosses): The losses object from the train step """ + pass def transform_gradients(self, losses: TrainingLosses) -> None: diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index aa492444c..e89a86d50 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -940,9 +940,6 @@ def update_before_train(self, current_server_round: int) -> None: gc.freeze() def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: - nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm) - - def transform_gradients(self, losses: TrainingLosses) -> None: """ Apply the gradient clipping performed by the default nnunet trainer. This is the default behavior for nnunet 2.5.1 @@ -950,4 +947,4 @@ def transform_gradients(self, losses: TrainingLosses) -> None: Args: losses (TrainingLosses): Not used for this transformation. """ - self._transform_gradients(self.model, losses) + nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm) From 33a35f633843f0e73046b6f37977b80be7440a2d Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:46:12 -0400 Subject: [PATCH 25/32] update README of nnunet example --- examples/nnunet_pfl_example/README.md | 64 +++------------------------ 1 file changed, 7 insertions(+), 57 deletions(-) diff --git a/examples/nnunet_pfl_example/README.md b/examples/nnunet_pfl_example/README.md index e5f890955..159b21beb 100644 --- a/examples/nnunet_pfl_example/README.md +++ b/examples/nnunet_pfl_example/README.md @@ -1,69 +1,19 @@ -# NnUNetClient Example +# NnUNetClient With Personalization Example -This example demonstrates how to use the NnunetClient to train nnunet segmentation models in a federated setting. - -By default this example trains an nnunet model on the Task04_Hippocampus dataset from the Medical Segmentation Decathlon (MSD). However, any of the MSD datasets can be used by specifying them with the msd_dataset_name flag for the client. To run this example first create a config file for the server. An example config has been provided in this directory. The required keys for the config are: - -```yaml -# Parameters that describe the server -n_server_rounds: 1 - -# Parameters that describe the clients -n_clients: 1 -local_epochs: 1 # Or local_steps, one or the other must be chosen - -nnunet_config: 2d -``` - -The only additional parameter required by nnunet is nnunet_config which is one of the official nnunet configurations (2d, 3d_fullres, 3d_lowres, 3d_cascade_fullres) - -One may also add the following optional keys to the config yaml file. If a nnunet plans file (which specifies model architecture and training hyperparameters) is not provided, the server will ask one of the clients to generate one using nnunet. - -```yaml -# Optional config parameters -nnunet_plans: /Path/to/nnunet/plans.json -starting_checkpoint: /Path/to/starting/checkpoint.pth -``` +Building on the [nnunet_example](../nnunet_example/README.md), here we demonstrate how personalized +methods can be applied to `NnUNetClient` class. The config requirements remain the same as in the original +`nnunet_example`. To run a federated learning experiment with nnunet models, first ensure you are in the FL4Health directory and then start the nnunet server using the following command. To view a list of optional flags use the --help flag ```bash -python -m examples.nnunet_example.server --config_path examples/nnunet_example/config.yaml +python -m examples.nnunet_pfl_example.server --config_path examples/nnunet_pfl_example/config.yaml ``` Once the server has started, start the necessary number of clients specified by the n_clients key in the config file. Each client can be started by running the following command in a separate session. To view a list of optional flags use the --help flag. ```bash -python -m examples.nnunet_example.client --dataset_path examples/datasets/nnunet +python -m examples.nnunet_pfl_example.client --dataset_path examples/datasets/nnunet --personalized-strategy ditto ``` -The MSD dataset will be downloaded and prepared automatically by the nnunet example script if it does not already exist. The dataset_path flag is used as more of a data working directory by the client. The client will create nnunet_raw, nnunet_preprocessed and nnunet_results sub directories if they do not already exist in the dataset_path folder. The dataset itself will be stored in a folder within nnunet_raw. Therefore when checking if the data already exists, the client will look for the following folder '{dataset_path}/nnunet_raw/{dataset_name}' - -## Definitions - -### Logits: - -The outputs of a model prior to the activation function. Values are unconstrained (-inf, inf) - -### Probabilities: - -Probabilities are values constrained to the range (0, 1) that represent the models confidence of a pixel/voxel being part of a particular class. Similar to other DNN's, they do not necessarily represent actual probabilities, however it is sometimes convenient to interpret them as such. The predicted probabilities are the outputs of a model after a normalizing activation function such as softmax or sigmoid. A 2d example is shown below. In some instances one might have labels that are probabilities as opposed to integer class labels in which case we would refer to them as ground-truth probabilities. - -### Detection Maps: - -Images that contain an arbitrary number of distinct detected volumes generally derived from the predicted probabilities of a segmentation model. Values are constrained to range [0, 1]. Detected volumes are defined as: -- Each detected volume is a connected component that must be non-connected and non-overlapping (mutually exclusive) with other volumes of the same class. (Therefore detection maps for multiclass segmentation must be one hot encoded) -- Each pixel/voxel within a volume must have the same predicted probability. Therefore there is a single confidence/likelihood score for each volume. - -Detected volumes typically also have a minimum size determined by the number of pixels/voxels that are a part of the volume. Detection maps may be computed from probabilities in a variety of ways. One example used for 3d medical images can be found in the [report guided annotation](https://github.com/DIAGNijmegen/Report-Guided-Annotation) API. An example of a 2d detection map is shown below. - -### Segmentations: - -Images in which pixels have been labelled or assigned one or more specific integer classes. If one hot encoded they must be binary {0, 1} or boolean {False, True} tensors. If not one hot encoded they must be be tensors containing only integers that represent the class labels (eg. constrained to {0, 1, 2, ..., N}). The labels/targets for segmentation models may be referred to as ground-truth segmentations. Predicted segmentations refers to model outputs (which are usually probability maps) that have been processed or thresholded in some way to adhere to the definition of a segmentation. They usually represents the model's final prediction of the class with no information on confidence. An example of a binary or one-hot-encoded predicted segmentation is shown below. - -#### Examples of different output types in the case of 2d binary segmentation: - - - | | | -:----------------------------:|:------------------------------:|:----------------------------------: -Probabilities | Detection Map | Segmentation| +The same MSD dataset that was used in the original `nnunet_example` is also used here. From 026cb036235739883560b9cf84f2af6e46965820 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:51:49 -0400 Subject: [PATCH 26/32] add _transform_gradients to protocol --- fl4health/mixins/core_protocols.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index 5d21f23f7..1950f80af 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -106,6 +106,9 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp def transform_target(self, target: TorchTargetType) -> TorchTargetType: pass # pragma: no cover + def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: + pass # pragma: no cover + def transform_gradients(self, losses: TrainingLosses) -> None: pass # pragma: no cover From ead64735fabeac56d5f882c01ae0007ed30d7aee Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 02:53:18 -0400 Subject: [PATCH 27/32] rename _transform_gradients_with_model --- fl4health/clients/basic_client.py | 4 ++-- fl4health/clients/nnunet_client.py | 4 ++-- fl4health/mixins/core_protocols.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py index 9cb096e5a..258a7c2a8 100644 --- a/fl4health/clients/basic_client.py +++ b/fl4health/clients/basic_client.py @@ -1389,7 +1389,7 @@ def update_before_epoch(self, epoch: int) -> None: """ pass - def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: + def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None: """Helper transform gradients method that allows for injection of model. NOTE: Subclasses should implement this helper should there be a need to specialize the logic @@ -1410,7 +1410,7 @@ def transform_gradients(self, losses: TrainingLosses) -> None: Args: losses (TrainingLosses): The losses object from the train step """ - return self._transform_gradients(self.model, losses) + return self._transform_gradients_with_model(self.model, losses) def _save_client_state(self) -> None: """ diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index e89a86d50..8aea0c8ac 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -232,7 +232,7 @@ def _apply_backwards_on_losses_and_take_step( # Rescale gradients then clip based on specified norm self.grad_scaler.unscale_(optimizer) - self._transform_gradients(model, losses) + self._transform_gradients_with_model(model, losses) # Update parameters and scaler self.grad_scaler.step(optimizer) @@ -939,7 +939,7 @@ 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, model: torch.nn.Module, losses: TrainingLosses) -> None: + def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None: """ Apply the gradient clipping performed by the default nnunet trainer. This is the default behavior for nnunet 2.5.1 diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py index 1950f80af..6d044edcb 100644 --- a/fl4health/mixins/core_protocols.py +++ b/fl4health/mixins/core_protocols.py @@ -106,7 +106,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp def transform_target(self, target: TorchTargetType) -> TorchTargetType: pass # pragma: no cover - def _transform_gradients(self, model: torch.nn.Module, losses: TrainingLosses) -> None: + def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None: pass # pragma: no cover def transform_gradients(self, losses: TrainingLosses) -> None: From b77e7a245ef85572ac4221ba0fab74f10e940265 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 09:34:07 -0400 Subject: [PATCH 28/32] cr --- fl4health/clients/nnunet_client.py | 75 ++++++++++++++++-------------- 1 file changed, 39 insertions(+), 36 deletions(-) diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py index 8aea0c8ac..ea474d2f3 100644 --- a/fl4health/clients/nnunet_client.py +++ b/fl4health/clients/nnunet_client.py @@ -208,6 +208,19 @@ def __init__( def _compute_preds_and_losses( self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType ) -> tuple[TrainingLosses, TorchPredType]: + """ + Overrides parent to include mixed precision training (autocasting and corresponding gradient scaling) as per + the original ``nnUNetTrainer``. + + Args: + input (TorchInputType): The input to be fed into the model. + target (TorchTargetType): The target corresponding to the input. + + Returns: + Tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any + predictions produced by the model. + """ + # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. if self.device.type != "cuda": return super()._compute_preds_and_losses(model, optimizer, input, target) @@ -226,6 +239,18 @@ def _compute_preds_and_losses( def _apply_backwards_on_losses_and_take_step( self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses ) -> TrainingLosses: + """ + Overrides parent to make use grad scaler with ``nnUNetTrainer``. + + Args: + model (nn.Module): the model used for making predictions. Passed here in case subclasses need it. + optimizer (Optimizer): the optimizer with which we take the step + losses (TrainingLosses): the losses to apply backwards on + + Returns: + TrainingLosses: The losses object after backwards application + """ + # Compute scaled loss and perform backward pass scaled_backward_loss = self.grad_scaler.scale(losses.backward["backward"]) scaled_backward_loss.backward() @@ -240,26 +265,6 @@ def _apply_backwards_on_losses_and_take_step( return losses - def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]: - """ - Given a single batch of input and target data, generate predictions, compute loss, update parameters and - optionally update metrics if they exist. (i.e. backprop on a single batch of data). Assumes ``self.model`` is - in train mode already. - - Overrides parent to include mixed precision training (autocasting and corresponding gradient scaling) as per - the original ``nnUNetTrainer``. - - Args: - input (TorchInputType): The input to be fed into the model. - target (TorchTargetType): The target corresponding to the input. - - Returns: - Tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any - predictions produced by the model. - """ - # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method. - return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target) - @use_default_signal_handlers # Dataloaders use multiprocessing def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]: """ @@ -620,6 +625,20 @@ def setup_client(self, config: Config) -> None: def _predict_with_model( self, model: torch.nn.Module, input: TorchInputType ) -> tuple[TorchPredType, dict[str, torch.Tensor]]: + """ + Generate model outputs. Overridden because nnunets outputs lists when deep supervision is on so we have to + reformat the output into dicts. + + Additionally if device type is cuda, loss computed in mixed precision. + + Args: + model (nn.Module): The model used to make predictions + input (TorchInputType): The model inputs + + Returns: + tuple[TorchPredType, dict[str, torch.Tensor]]: A tuple in which the first element model outputs indexed by + name. The second element is unused by this subclass and therefore is always an empty dict + """ if isinstance(input, torch.Tensor): # If device type is cuda, nnUNet defaults to mixed precision forward pass if self.device.type == "cuda": @@ -642,22 +661,6 @@ def _predict_with_model( "Was expecting nnunet model output to be either a torch.Tensor or a list/tuple of torch.Tensors" ) - def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]: - """ - Generate model outputs. Overridden because nnunets outputs lists when deep supervision is on so we have to - reformat the output into dicts. - - Additionally if device type is cuda, loss computed in mixed precision. - - Args: - input (TorchInputType): The model inputs - - Returns: - tuple[TorchPredType, dict[str, torch.Tensor]]: A tuple in which the first element model outputs indexed by - name. The second element is unused by this subclass and therefore is always an empty dict - """ - return self._predict_with_model(self.model, input) - def compute_loss_and_additional_losses( self, preds: TorchPredType, From 34d0bcc465e2e84ff19414aefeb09e9f289a9115 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 10:05:42 -0400 Subject: [PATCH 29/32] test predict ditto mixin --- tests/mixins/personalized/test_ditto.py | 71 ++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py index 76a04e41b..eb1efe2c5 100644 --- a/tests/mixins/personalized/test_ditto.py +++ b/tests/mixins/personalized/test_ditto.py @@ -28,8 +28,7 @@ from fl4health.parameter_exchange.parameter_packer import ( ParameterPackerAdaptiveConstraint, ) - -# from fl4health.utils.losses import TrainingLosses +from fl4health.utils.losses import TrainingLosses from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType @@ -157,6 +156,74 @@ def test_get_parameters() -> None: assert_array_equal(packed_params, pack_params_return_val) +@patch.object(_TestDittoedClient, "compute_penalty_loss") +@patch.object(_TestDittoedClient, "_apply_backwards_on_losses_and_take_step") +@patch.object(_TestDittoedClient, "_compute_preds_and_losses") +def test_predict( + mock_private_compute_preds_and_losses: MagicMock, + mock_private_apply_backwards_on_losses_and_take_step: MagicMock, + mock_compute_penalty_loss: MagicMock, +) -> None: + # setup client + client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) + client.model = torch.nn.Linear(5, 5) + client.global_model = torch.nn.Linear(5, 5) + client.optimizers = { + "global": torch.optim.SGD(client.model.parameters(), lr=0.0001), + "local": torch.optim.SGD(client.model.parameters(), lr=0.0001), + } + # setup mocks + mock_param_exchanger = MagicMock() + push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float) + pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float) + mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights + mock_param_exchanger.pack_parameters.return_value = pack_params_return_val + client.parameter_exchanger = mock_param_exchanger + client.initialized = True + + mock_backward_loss = MagicMock() + mock_additional_global_loss = MagicMock() + dummy_training_losses = TrainingLosses( + backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss} + ) + dummy_training_losses_for_local = TrainingLosses( + backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss} + ) + mock_private_compute_preds_and_losses.side_effect = [ + ( + dummy_training_losses, + {"prediction": torch.Tensor([1, 2, 3, 4, 5])}, + ), + ( + dummy_training_losses_for_local, + {"prediction": torch.Tensor([1, 2, 3, 4, 5])}, + ), + ] + mock_private_apply_backwards_on_losses_and_take_step.side_effect = [ + dummy_training_losses, + dummy_training_losses_for_local, + ] + + # act + input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3) + # TODO: fix the mixin/protocol typing that leads to mypy complaint + _ = client.train_step(input, target) # type: ignore + + mock_private_compute_preds_and_losses.assert_has_calls( + [ + _Call(((client.global_model, client.optimizers["global"], input, target), {})), + _Call(((client.model, client.optimizers["local"], input, target), {})), + ] + ) + mock_private_apply_backwards_on_losses_and_take_step.assert_has_calls( + [ + _Call(((client.global_model, client.optimizers["global"], dummy_training_losses), {})), + _Call(((client.model, client.optimizers["local"], dummy_training_losses_for_local), {})), + ] + ) + mock_compute_penalty_loss.assert_called_once() + + @patch.object(_TestDittoedClient, "setup_client") @patch("fl4health.mixins.personalized.ditto.FullParameterExchanger") def test_get_parameters_uninitialized(mock_param_exchanger: MagicMock, mock_setup_client: MagicMock) -> None: From ab7f7a0b76e713215b2d2abd85fecef91b603114 Mon Sep 17 00:00:00 2001 From: Val Andrei Fajardo Date: Fri, 30 May 2025 10:11:06 -0400 Subject: [PATCH 30/32] test val step --- tests/mixins/personalized/test_ditto.py | 53 ++++++++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py index eb1efe2c5..52f4a59a3 100644 --- a/tests/mixins/personalized/test_ditto.py +++ b/tests/mixins/personalized/test_ditto.py @@ -28,7 +28,7 @@ from fl4health.parameter_exchange.parameter_packer import ( ParameterPackerAdaptiveConstraint, ) -from fl4health.utils.losses import TrainingLosses +from fl4health.utils.losses import EvaluationLosses, TrainingLosses from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType @@ -333,3 +333,54 @@ def test_safe_model_raises_error() -> None: with pytest.raises(ValueError): # TODO: fix the mixin/protocol typing that leads to mypy complaint client.safe_global_model() # type: ignore + + +@patch.object(_TestDittoedClient, "_val_step_with_model") +def test_val_step( + mock_private_val_step_with_model: MagicMock, +) -> None: + # setup client + client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu")) + client.model = torch.nn.Linear(5, 5) + client.global_model = torch.nn.Linear(5, 5) + client.optimizers = { + "global": torch.optim.SGD(client.model.parameters(), lr=0.0001), + "local": torch.optim.SGD(client.model.parameters(), lr=0.0001), + } + # setup mocks + mock_param_exchanger = MagicMock() + push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float) + pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float) + mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights + mock_param_exchanger.pack_parameters.return_value = pack_params_return_val + client.parameter_exchanger = mock_param_exchanger + client.initialized = True + + dummy_training_losses = EvaluationLosses( + checkpoint=torch.ones(5), + ) + dummy_training_losses_for_local = EvaluationLosses( + checkpoint=torch.ones(5), + ) + mock_private_val_step_with_model.side_effect = [ + ( + dummy_training_losses, + {"prediction": torch.Tensor([1, 2, 3, 4, 5])}, + ), + ( + dummy_training_losses_for_local, + {"prediction": torch.Tensor([1, 2, 3, 4, 5])}, + ), + ] + + # act + input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3) + # TODO: fix the mixin/protocol typing that leads to mypy complaint + _ = client.val_step(input, target) # type: ignore + + mock_private_val_step_with_model.assert_has_calls( + [ + _Call(((client.global_model, input, target), {})), + _Call(((client.model, input, target), {})), + ] + ) From 4e4ab2887c68f77beb7bd11816c26b93d7c582b9 Mon Sep 17 00:00:00 2001 From: David Emerson <43939939+emersodb@users.noreply.github.com> Date: Thu, 5 Jun 2025 09:14:18 -0400 Subject: [PATCH 31/32] Typo fix --- fl4health/mixins/personalized/ditto.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index 42cc0a8c3..a25c86a76 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -127,7 +127,7 @@ def _copy_optimizer_with_new_params(self: DittoPersonalizedProtocol, original_op """Helper method to make a copy of the original optimizer for the global model. Args: - original_optimizer (Optimizer): original optimizer of the underyling `BasicClient`. + original_optimizer (Optimizer): original optimizer of the underlying `BasicClient`. Returns: Optimizer: a copy of the original optimizer to be used by the global model. From 51208055ff2267691a10c4014bf3853adf4ef302 Mon Sep 17 00:00:00 2001 From: David Emerson <43939939+emersodb@users.noreply.github.com> Date: Thu, 5 Jun 2025 09:14:47 -0400 Subject: [PATCH 32/32] Formatting --- fl4health/mixins/personalized/ditto.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py index a25c86a76..6137aa81c 100644 --- a/fl4health/mixins/personalized/ditto.py +++ b/fl4health/mixins/personalized/ditto.py @@ -95,10 +95,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None: f"base classes implement BasicClient. This may cause runtime errors." ) log(WARN, msg) - warnings.warn( - msg, - RuntimeWarning, - ) + warnings.warn(msg, RuntimeWarning) def safe_global_model(self: DittoPersonalizedProtocol) -> nn.Module: """Convenient accessor for the global model.