diff --git a/monai/fl/client/monai_algo.py b/monai/fl/client/monai_algo.py index 175ed05d67..013fe0ed1b 100644 --- a/monai/fl/client/monai_algo.py +++ b/monai/fl/client/monai_algo.py @@ -18,7 +18,7 @@ import monai from monai.bundle import ConfigParser -from monai.bundle.config_item import ConfigItem +from monai.bundle.config_item import ConfigComponent, ConfigItem from monai.config import IgniteInfo from monai.fl.client.client_algo import ClientAlgo from monai.fl.utils.constants import ( @@ -27,6 +27,7 @@ FiltersType, FlPhase, FlStatistics, + ModelType, RequiredBundleKeys, WeightType, ) @@ -81,6 +82,14 @@ def check_bundle_config(parser): raise KeyError(f"Bundle config misses required key `{k}`") +def remove_ckpt_loader(parser): + if BundleKeys.VALIDATE_HANDLERS in parser: + for h in parser[BundleKeys.VALIDATE_HANDLERS]: + if ConfigComponent.is_instantiable(h): + if "CheckpointLoader" in h["_target_"]: + h["_disabled_"] = True + + class MonaiAlgo(ClientAlgo): """ Implementation of ``ClientAlgo`` to allow federated learning with MONAI bundle configurations. @@ -89,9 +98,17 @@ class MonaiAlgo(ClientAlgo): bundle_root: path of bundle. local_epochs: number of local epochs to execute during each round of local training; defaults to 1. send_weight_diff: whether to send weight differences rather than full weights; defaults to `True`. - config_train_filename: bundle training config path relative to bundle_root; defaults to "configs/train.json" - config_evaluate_filename: bundle evaluation config path relative to bundle_root; defaults to "configs/evaluate.json" + config_train_filename: bundle training config path relative to bundle_root; defaults to "configs/train.json". + config_evaluate_filename: bundle evaluation config path relative to bundle_root; defaults to "configs/evaluate.json". config_filters_filename: filter configuration file. + disable_ckpt_loader: do not use CheckpointLoader if defined in train/evaluate configs; defaults to `True`. + best_model_filepath: location of best model checkpoint; defaults "models/model.pt" relative to `bundle_root`. + final_model_filepath: location of final model checkpoint; defaults "models/model_final.pt" relative to `bundle_root`. + seed: set random seed for modules to enable or disable deterministic training; defaults to `None`, + i.e., non-deterministic training. + benchmark: set benchmark to `False` for full deterministic behavior in cuDNN components. + Note, full determinism in federated learning depends also on deterministic behavior of other FL components, + e.g., the aggregator, which is not controlled by this class. """ def __init__( @@ -102,6 +119,11 @@ def __init__( config_train_filename: Optional[str] = "configs/train.json", config_evaluate_filename: Optional[str] = "configs/evaluate.json", config_filters_filename: Optional[str] = None, + disable_ckpt_loader: bool = True, + best_model_filepath: Optional[str] = "models/model.pt", + final_model_filepath: Optional[str] = "models/model_final.pt", + seed: Optional[int] = None, + benchmark: bool = True, ): self.logger = logging.getLogger(self.__class__.__name__) @@ -111,6 +133,10 @@ def __init__( self.config_train_filename = config_train_filename self.config_evaluate_filename = config_evaluate_filename self.config_filters_filename = config_filters_filename + self.disable_ckpt_loader = disable_ckpt_loader + self.model_filepaths = {ModelType.BEST_MODEL: best_model_filepath, ModelType.FINAL_MODEL: final_model_filepath} + self.seed = seed + self.benchmark = benchmark self.app_root = None self.train_parser = None @@ -133,13 +159,17 @@ def initialize(self, extra=None): Args: extra: Dict with additional information that should be provided by FL system, - i.e., `ExtraItems.CLIENT_NAME` and `ExtraItems.APP_ROOT`. + i.e., `ExtraItems.CLIENT_NAME` and `ExtraItems.APP_ROOT`. """ if extra is None: extra = {} - self.logger.info(f"Initializing {self.client_name} ...") self.client_name = extra.get(ExtraItems.CLIENT_NAME, "noname") + self.logger.info(f"Initializing {self.client_name} ...") + + if self.seed: + monai.utils.set_determinism(seed=self.seed) + torch.backends.cudnn.benchmark = self.benchmark # FL platform needs to provide filepath to configuration files self.app_root = extra.get(ExtraItems.APP_ROOT, "") @@ -178,9 +208,10 @@ def initialize(self, extra=None): if BundleKeys.TRAIN_TRAINER_MAX_EPOCHS in self.train_parser: self.train_parser[BundleKeys.TRAIN_TRAINER_MAX_EPOCHS] = self.local_epochs - self.train_parser.parse() - self.eval_parser.parse() - self.filter_parser.parse() + # remove checkpoint loaders + if self.disable_ckpt_loader: + remove_ckpt_loader(self.train_parser) + remove_ckpt_loader(self.eval_parser) # Get trainer, evaluator self.trainer = self.train_parser.get_parsed_content( @@ -233,7 +264,9 @@ def train(self, data: ExchangeObject, extra=None): # get current iteration when a round starts self.iter_of_start_time = self.trainer.state.iteration - copy_model_state(src=self.global_weights, dst=self.trainer.network) + _, updated_keys, _ = copy_model_state(src=self.global_weights, dst=self.trainer.network) + if len(updated_keys) == 0: + self.logger.warning("No weights loaded!") self.logger.info(f"Start {self.client_name} training...") self.trainer.run() @@ -245,26 +278,52 @@ def get_weights(self, extra=None): extra: Dict with additional information that can be provided by FL system. Returns: - return_weights: `ExchangeObject` containing current weights. + return_weights: `ExchangeObject` containing current weights (default) + or load requested model type from disk (`ModelType.BEST_MODEL` or `ModelType.FINAL_MODEL`). """ if extra is None: extra = {} + + # by default return current weights, return best if requested via model type. self.phase = FlPhase.GET_WEIGHTS - if self.trainer: - weights = get_state_dict(self.trainer.network) - weigh_type = WeightType.WEIGHTS - stats = self.trainer.get_stats() - # calculate current iteration and epoch data after training. - stats[FlStatistics.NUM_EXECUTED_ITERATIONS] = self.trainer.state.iteration - self.iter_of_start_time - # compute weight differences - if self.send_weight_diff: - weights = compute_weight_diff(global_weights=self.global_weights, local_var_dict=weights) - weigh_type = WeightType.WEIGHT_DIFF + + if ExtraItems.MODEL_TYPE in extra: + model_type = extra.get(ExtraItems.MODEL_TYPE) + if not isinstance(model_type, ModelType): + raise ValueError( + f"Expected requested model type to be of type `ModelType` but received {type(model_type)}" + ) + if model_type in self.model_filepaths: + model_path = os.path.join(self.bundle_root, self.model_filepaths[model_type]) + if not os.path.isfile(model_path): + raise ValueError(f"No best model checkpoint exists at {model_path}") + weights = torch.load(model_path, map_location="cpu") + weigh_type = WeightType.WEIGHTS + stats = dict() + self.logger.info(f"Returning best checkpoint weights from {model_path}.") + else: + raise ValueError( + f"Requested model type {model_type} not specified in `model_filepahts`: {self.model_filepaths}" + ) else: - weights = None - weigh_type = None - stats = dict() + if self.trainer: + weights = get_state_dict(self.trainer.network) + weigh_type = WeightType.WEIGHTS + stats = self.trainer.get_stats() + # calculate current iteration and epoch data after training. + stats[FlStatistics.NUM_EXECUTED_ITERATIONS] = self.trainer.state.iteration - self.iter_of_start_time + # compute weight differences + if self.send_weight_diff: + weights = compute_weight_diff(global_weights=self.global_weights, local_var_dict=weights) + weigh_type = WeightType.WEIGHT_DIFF + self.logger.info("Returning current weight differences.") + else: + self.logger.info("Returning current weights.") + else: + weights = None + weigh_type = None + stats = dict() if not isinstance(stats, dict): raise ValueError(f"stats is not a dict, {stats}") @@ -311,7 +370,9 @@ def evaluate(self, data: ExchangeObject, extra=None): global_weights=data.weights, local_var_dict=get_state_dict(self.evaluator.network) ) - copy_model_state(src=global_weights, dst=self.evaluator.network) + _, updated_keys, _ = copy_model_state(src=global_weights, dst=self.evaluator.network) + if len(updated_keys) == 0: + self.logger.warning("No weights loaded!") self.logger.info(f"Start {self.client_name} evaluating...") if isinstance(self.trainer, monai.engines.Trainer): self.evaluator.run(self.trainer.state.epoch + 1) diff --git a/monai/fl/utils/constants.py b/monai/fl/utils/constants.py index 90b452e70d..f6da8d4ea0 100644 --- a/monai/fl/utils/constants.py +++ b/monai/fl/utils/constants.py @@ -17,9 +17,14 @@ class WeightType(StrEnum): WEIGHT_DIFF = "fl_weight_diff" +class ModelType(StrEnum): + BEST_MODEL = "fl_best_model" + FINAL_MODEL = "fl_final_model" + + class ExtraItems(StrEnum): ABORT = "fl_abort" - MODEL_NAME = "fl_model_name" + MODEL_TYPE = "fl_model_type" CLIENT_NAME = "fl_client_name" APP_ROOT = "fl_app_root" @@ -43,6 +48,7 @@ class BundleKeys(StrEnum): TRAINER = "train#trainer" EVALUATOR = "validate#evaluator" TRAIN_TRAINER_MAX_EPOCHS = "train#trainer#max_epochs" + VALIDATE_HANDLERS = "validate#handlers" class FiltersType(StrEnum):