update basic client with helper methods

Author: nerdai

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:
         """

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)

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
 

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)