add FOR_GLOBAL_MODEL_KEY in config.utils

Author: nerdai

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:

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

fl4health/utils/config.py

@@ -8,6 +8,8 @@
     "batch_size": int,
 }
 
+FOR_GLOBAL_MODEL_KEY = "__for_global_model"
+
 T = TypeVar("T")