Extend supervised train tests

Author: C-Achard

napari_cellseg3d/_tests/fixtures.py

@@ -36,23 +36,42 @@ def forward(self, x):
         return self.forward_encoder(x), self.forward_decoder(x)
 
 
-class OptimizerFixture:
-    def __call__(self, x):
+class ModelFixture(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mock_conv = torch.nn.Conv3d(1, 1, 1)
+        self.mock_conv.requires_grad_(False)
+
+    def forward(self, x):
         return x
 
+
+class OptimizerFixture:
+    def __init__(self):
+        self.param_groups = []
+        self.param_groups.append({"lr": 0})
+
     def zero_grad(self):
         pass
 
-    def step(self):
+    def step(self, *args):
+        pass
+
+
+class SchedulerFixture:
+    def step(self, *args):
         pass
 
 
 class LossFixture:
-    def __call__(self, x):
-        return x
+    def __call__(self, *args):
+        return self
 
-    def backward(self, x):
+    def backward(self, *args):
         pass
 
     def item(self):
         return 0
+
+    def detach(self):
+        return self

napari_cellseg3d/_tests/test_training.py

@@ -5,7 +5,9 @@
 from napari_cellseg3d._tests.fixtures import (
     LogFixture,
     LossFixture,
+    ModelFixture,
     OptimizerFixture,
+    SchedulerFixture,
     WNetFixture,
 )
 from napari_cellseg3d.code_models.models.model_test import TestModel
@@ -33,6 +35,7 @@ def test_supervised_training(make_napari_viewer_proxy):
     widget.labels_filepaths = [im_path_str]
     widget.epoch_choice.setValue(1)
     widget.val_interval_choice.setValue(1)
+    widget.device_choice.setCurrentIndex(0)
 
     assert widget.check_ready()
 
@@ -49,13 +52,19 @@ def test_supervised_training(make_napari_viewer_proxy):
     worker.config.val_data_dict = [
         {"image": im_path_str, "label": im_path_str}
     ]
-    worker.config.max_epochs = 1
+    worker.config.max_epochs = 2
     worker.config.validation_interval = 2
-    worker.log_parameters()
-    res = next(worker.train())
 
-    assert isinstance(res, TrainingReport)
-    assert res.epoch == 0
+    worker.log_parameters()
+    for res_i in worker.train(
+        provided_model=ModelFixture(),
+        provided_optimizer=OptimizerFixture(),
+        provided_loss=LossFixture(),
+        provided_scheduler=SchedulerFixture(),
+    ):
+        assert isinstance(res_i, TrainingReport)
+        res = res_i
+    assert res.epoch == 1
 
     widget.worker = worker
     res.show_plot = True
@@ -86,15 +95,15 @@ def test_unsupervised_training(make_napari_viewer_proxy):
         additional_results_description="wnet_test"
     )
     assert widget.worker.config.train_data_dict is not None
-    res = next(
-        widget.worker.train(
-            provided_model=WNetFixture(),
-            provided_optimizer=OptimizerFixture(),
-            provided_loss=LossFixture(),
-        )
-    )
-    assert isinstance(res, TrainingReport)
-    assert not res.show_plot
+    widget.worker.config.max_epochs = 1
+    for res_i in widget.worker.train(
+        provided_model=WNetFixture(),
+        provided_optimizer=OptimizerFixture(),
+        provided_loss=LossFixture(),
+    ):
+        assert isinstance(res_i, TrainingReport)
+        res = res_i
+    assert res.epoch == 0
     widget.worker._abort_requested = True
     res = next(
         widget.worker.train(

napari_cellseg3d/code_models/worker_training.py

@@ -999,7 +999,13 @@ def log_parameters(self):
         # self.log("\n")
         # self.log("-" * 20)
 
-    def train(self):
+    def train(
+        self,
+        provided_model=None,
+        provided_optimizer=None,
+        provided_loss=None,
+        provided_scheduler=None,
+    ):
         """Trains the PyTorch model for the given number of epochs, with the selected model and data,
         using the chosen batch size, validation interval, loss function, and number of samples.
         Will perform validation once every :py:obj:`val_interval` and save results if the mean dice is better
@@ -1070,13 +1076,16 @@ def train(self):
                     self.config.train_data_dict[0]
                 )
                 check = data_check["image"].shape
-
             do_sampling = self.config.sampling
-
             size = self.config.sample_size if do_sampling else check
-
             PADDING = utils.get_padding_dim(size)
-            model = model_class(input_img_size=PADDING, use_checkpoint=True)
+
+            model = (
+                model_class(input_img_size=PADDING, use_checkpoint=True)
+                if provided_model is None
+                else provided_model
+            )
+
             device = torch.device(self.config.device)
             model = model.to(device)
 
@@ -1276,8 +1285,10 @@ def get_patch_loader_func(num_samples):
             logger.info("\nDone")
 
             logger.debug("Optimizer")
-            optimizer = torch.optim.Adam(
-                model.parameters(), self.config.learning_rate
+            optimizer = (
+                torch.optim.Adam(model.parameters(), self.config.learning_rate)
+                if provided_optimizer is None
+                else provided_optimizer
             )
 
             factor = self.config.scheduler_factor
@@ -1286,12 +1297,16 @@ def get_patch_loader_func(num_samples):
                 self.log("Setting it to 0.5")
                 factor = 0.5
 
-            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-                optimizer=optimizer,
-                mode="min",
-                factor=factor,
-                patience=self.config.scheduler_patience,
-                verbose=VERBOSE_SCHEDULER,
+            scheduler = (
+                torch.optim.lr_scheduler.ReduceLROnPlateau(
+                    optimizer=optimizer,
+                    mode="min",
+                    factor=factor,
+                    patience=self.config.scheduler_patience,
+                    verbose=VERBOSE_SCHEDULER,
+                )
+                if provided_scheduler is None
+                else provided_scheduler
             )
             dice_metric = DiceMetric(
                 include_background=True, reduction="mean", ignore_empty=False
@@ -1342,6 +1357,8 @@ def get_patch_loader_func(num_samples):
 
             # device = torch.device(self.config.device)
             self.set_loss_from_config()
+            if provided_loss is not None:
+                self.loss_function = provided_loss
 
             # if model_name == "test":
             #     self.quit()