Fix loss function instantiation

Author: C-Achard

napari_cellseg3d/_tests/test_training.py

@@ -53,6 +53,14 @@ def test_update_loss_plot(make_napari_viewer_proxy):
     assert widget.train_loss_plot is not None
 
 
+def test_check_matching_losses():
+    plugin = Trainer(None)
+    config = plugin._set_worker_config()
+    worker = plugin._create_worker_from_config(config)
+
+    assert plugin.loss_list == list(worker.loss_dict.keys())
+
+
 def test_training(make_napari_viewer_proxy, qtbot):
     im_path = str(Path(__file__).resolve().parent / "res/test.tif")
 
@@ -73,9 +81,6 @@ def test_training(make_napari_viewer_proxy, qtbot):
 
     assert widget.check_ready()
 
-    #################
-    # Training is too long to test properly this way. Do not use on Github
-    #################
     MODEL_LIST["test"] = TestModel
     widget.model_choice.addItem("test")
     widget.model_choice.setCurrentText("test")

napari_cellseg3d/code_models/worker_training.py

@@ -15,6 +15,12 @@
     pad_list_data_collate,
 )
 from monai.inferers import sliding_window_inference
+from monai.losses import (
+    DiceCELoss,
+    DiceLoss,
+    GeneralizedDiceLoss,
+    TverskyLoss,
+)
 from monai.metrics import DiceMetric
 from monai.transforms import (
     # AsDiscrete,
@@ -125,6 +131,25 @@ def __init__(
         #######################################
         self.downloader = WeightsDownloader()
 
+        self.loss_dict = {
+            "Dice": DiceLoss(sigmoid=True),
+            # "BCELoss": torch.nn.BCELoss(), # dev
+            # "BCELogits": torch.nn.BCEWithLogitsLoss(),
+            "Generalized Dice": GeneralizedDiceLoss(sigmoid=True),
+            "DiceCE": DiceCELoss(sigmoid=True, lambda_ce=0.5),
+            "Tversky": TverskyLoss(sigmoid=True),
+            # "Focal loss": FocalLoss(),
+            # "Dice-Focal loss": DiceFocalLoss(sigmoid=True, lambda_dice=0.5),
+        }
+        self.loss_function = None
+
+    def set_loss_from_config(self):
+        try:
+            self.loss_function = self.loss_dict[self.config.loss_function]
+        except KeyError as e:
+            self.raise_error(e, "Loss function not found, aborting job")
+        return self.loss_function
+
     def set_download_log(self, widget):
         self.downloader.log_widget = widget
 
@@ -532,6 +557,7 @@ def get_loader_func(num_samples):
             self.log_parameters()
 
             device = torch.device(self.config.device)
+            self.set_loss_from_config()
 
             # if model_name == "test":
             #     self.quit()
@@ -571,7 +597,7 @@ def get_loader_func(num_samples):
                         ]  # TODO(cyril): adapt if additional channels
                         if len(outputs.shape) < 4:
                             outputs = outputs.unsqueeze(0)
-                    loss = self.config.loss_function(outputs, labels)
+                    loss = self.loss_function(outputs, labels)
                     loss.backward()
                     optimizer.step()
                     epoch_loss += loss.detach().item()

napari_cellseg3d/code_plugins/plugin_model_training.py

@@ -15,14 +15,6 @@
 if TYPE_CHECKING:
     import napari
 
-# MONAI
-from monai.losses import (
-    DiceCELoss,
-    DiceLoss,
-    GeneralizedDiceLoss,
-    TverskyLoss,
-)
-
 # Qt
 from qtpy.QtWidgets import QSizePolicy
 
@@ -139,17 +131,15 @@ def __init__(
         """Whether the worker should stop or not"""
         self.start_time = None
 
-        self.loss_dict = {
-            "Dice": DiceLoss(sigmoid=True),
-            # "BCELoss": torch.nn.BCELoss(), # dev
-            # "BCELogits": torch.nn.BCEWithLogitsLoss(),
-            "Generalized Dice": GeneralizedDiceLoss(sigmoid=True),
-            "DiceCE": DiceCELoss(sigmoid=True, lambda_ce=0.5),
-            "Tversky": TverskyLoss(sigmoid=True),
-            # "Focal loss": FocalLoss(),
-            # "Dice-Focal loss": DiceFocalLoss(sigmoid=True, lambda_dice=0.5),
-        }
-        """Dict of loss functions"""
+        self.loss_list = [  # MUST BE MATCHED WITH THE LOSS FUNCTIONS IN THE TRAINING WORKER DICT
+            "Dice",
+            "Generalized Dice",
+            "DiceCE",
+            "Tversky",
+            # "Focal loss",
+            # "Dice-Focal loss",
+        ]
+        """List of loss functions"""
 
         self.canvas = None
         """Canvas to plot loss and dice metric in"""
@@ -192,10 +182,7 @@ def __init__(
         )
 
         self.loss_choice = ui.DropdownMenu(
-            # sorted(
-            list(
-                self.loss_dict.keys(),
-            ),
+            self.loss_list,
             text_label="Loss function",
         )
         self.lbl_loss_choice = self.loss_choice.label
@@ -383,10 +370,6 @@ def _update_validation_choice(self):
         elif validation.maximum() < max_epoch:
             validation.setMaximum(max_epoch)
 
-    def get_loss(self, key):
-        """Getter for loss function selected by user"""
-        return self.loss_dict[key]
-
     def _toggle_patch_dims(self):
         if self.patch_choice.isChecked():
             [w.setVisible(True) for w in self.patch_size_widgets]
@@ -899,7 +882,7 @@ def _set_worker_config(self) -> config.TrainingWorkerConfig:
             train_data_dict=self.data,
             validation_percent=validation_percent,
             max_epochs=self.epoch_choice.value(),
-            loss_function=self.get_loss(self.loss_choice.currentText()),
+            loss_function=self.loss_choice.currentText(),
             learning_rate=float(self.learning_rate_choice.currentText()),
             scheduler_patience=self.scheduler_patience_choice.value(),
             scheduler_factor=self.scheduler_factor_choice.slider_value,