wip

Author: nerdai

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