Test unsupervised training and raise coverage

Author: C-Achard

napari_cellseg3d/_tests/fixtures.py

@@ -1,3 +1,4 @@
+import torch
 from qtpy.QtWidgets import QTextEdit
 
 from napari_cellseg3d.utils import LOGGER as logger
@@ -17,3 +18,41 @@ def warn(self, warning):
 
     def error(self, e):
         raise (e)
+
+
+class WNetFixture(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_encoder(self, x):
+        return x
+
+    def forward_decoder(self, x):
+        return x
+
+    def forward(self, x):
+        return self.forward_encoder(x), self.forward_decoder(x)
+
+
+class OptimizerFixture:
+    def __call__(self, x):
+        return x
+
+    def zero_grad(self):
+        pass
+
+    def step(self):
+        pass
+
+
+class LossFixture:
+    def __call__(self, x):
+        return x
+
+    def backward(self, x):
+        pass
+
+    def item(self):
+        return 0

…cellseg3d/_tests/test_training_plugin.py …cellseg3d/_tests/test_plugin_training.pynapari_cellseg3d/_tests/test_training_plugin.py renamed to napari_cellseg3d/_tests/test_plugin_training.py napari_cellseg3d/_tests/test_training_plugin.py renamed to napari_cellseg3d/_tests/test_plugin_training.py

@@ -1,9 +1,6 @@
 from pathlib import Path
 
 from napari_cellseg3d import config
-from napari_cellseg3d._tests.fixtures import LogFixture
-from napari_cellseg3d.code_models.models.model_test import TestModel
-from napari_cellseg3d.code_models.workers_utils import TrainingReport
 from napari_cellseg3d.code_plugins.plugin_model_training import (
     Trainer,
 )
@@ -109,51 +106,3 @@ def test_check_matching_losses():
     worker = plugin._create_supervised_worker_from_config(config)
 
     assert plugin.loss_list == list(worker.loss_dict.keys())
-
-
-def test_training(make_napari_viewer_proxy, qtbot):
-    viewer = make_napari_viewer_proxy()
-    widget = Trainer(viewer)
-    widget.log = LogFixture()
-    viewer.window.add_dock_widget(widget)
-
-    widget.images_filepath = []
-    widget.labels_filepaths = []
-
-    assert not widget.check_ready()
-
-    widget.images_filepaths = [im_path_str]
-    widget.labels_filepaths = [im_path_str]
-    widget.epoch_choice.setValue(1)
-    widget.val_interval_choice.setValue(1)
-
-    assert widget.check_ready()
-
-    MODEL_LIST["test"] = TestModel
-    widget.model_choice.addItem("test")
-    widget.model_choice.setCurrentText("test")
-    widget.unsupervised_mode = False
-    worker_config = widget._set_worker_config()
-    assert worker_config.model_info.name == "test"
-    worker = widget._create_supervised_worker_from_config(worker_config)
-    worker.config.train_data_dict = [
-        {"image": im_path_str, "label": im_path_str}
-    ]
-    worker.config.val_data_dict = [
-        {"image": im_path_str, "label": im_path_str}
-    ]
-    worker.config.max_epochs = 1
-    worker.config.validation_interval = 2
-    worker.log_parameters()
-    res = next(worker.train())
-
-    assert isinstance(res, TrainingReport)
-    assert res.epoch == 0
-
-    widget.worker = worker
-    res.show_plot = True
-    res.loss_1_values = {"loss": [1, 1, 1, 1]}
-    res.loss_2_values = [1, 1, 1, 1]
-    widget.on_yield(res)
-    assert widget.loss_1_values["loss"] == [1, 1, 1, 1]
-    assert widget.loss_2_values == [1, 1, 1, 1]

napari_cellseg3d/_tests/test_training.py

@@ -0,0 +1,94 @@
+from pathlib import Path
+
+from napari_cellseg3d._tests.fixtures import (
+    LogFixture,
+    LossFixture,
+    OptimizerFixture,
+    WNetFixture,
+)
+from napari_cellseg3d.code_models.models.model_test import TestModel
+from napari_cellseg3d.code_models.workers_utils import TrainingReport
+from napari_cellseg3d.code_plugins.plugin_model_training import (
+    Trainer,
+)
+from napari_cellseg3d.config import MODEL_LIST
+
+im_path = Path(__file__).resolve().parent / "res/test.tif"
+im_path_str = str(im_path)
+
+
+def test_supervised_training(make_napari_viewer_proxy):
+    viewer = make_napari_viewer_proxy()
+    widget = Trainer(viewer)
+    widget.log = LogFixture()
+
+    widget.images_filepath = []
+    widget.labels_filepaths = []
+
+    assert not widget.check_ready()
+
+    widget.images_filepaths = [im_path_str]
+    widget.labels_filepaths = [im_path_str]
+    widget.epoch_choice.setValue(1)
+    widget.val_interval_choice.setValue(1)
+
+    assert widget.check_ready()
+
+    MODEL_LIST["test"] = TestModel
+    widget.model_choice.addItem("test")
+    widget.model_choice.setCurrentText("test")
+    widget.unsupervised_mode = False
+    worker_config = widget._set_worker_config()
+    assert worker_config.model_info.name == "test"
+    worker = widget._create_supervised_worker_from_config(worker_config)
+    worker.config.train_data_dict = [
+        {"image": im_path_str, "label": im_path_str}
+    ]
+    worker.config.val_data_dict = [
+        {"image": im_path_str, "label": im_path_str}
+    ]
+    worker.config.max_epochs = 1
+    worker.config.validation_interval = 2
+    worker.log_parameters()
+    res = next(worker.train())
+
+    assert isinstance(res, TrainingReport)
+    assert res.epoch == 0
+
+    widget.worker = worker
+    res.show_plot = True
+    res.loss_1_values = {"loss": [1, 1, 1, 1]}
+    res.loss_2_values = [1, 1, 1, 1]
+    widget.on_yield(res)
+    assert widget.loss_1_values["loss"] == [1, 1, 1, 1]
+    assert widget.loss_2_values == [1, 1, 1, 1]
+
+
+def test_unsupervised_training(make_napari_viewer_proxy):
+    viewer = make_napari_viewer_proxy()
+    widget = Trainer(viewer)
+    widget.log = LogFixture()
+    widget.worker = None
+    widget._toggle_unsupervised_mode(enabled=True)
+    widget.model_choice.setCurrentText("WNet")
+
+    widget.patch_choice.setChecked(True)
+    [w.setValue(4) for w in widget.patch_size_widgets]
+
+    widget.unsupervised_images_filewidget.text_field.setText(
+        str(im_path.parent)
+    )
+    # widget.start()
+    widget.data = widget.create_dataset_dict_no_labs()
+    widget.worker = widget._create_worker(
+        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)

napari_cellseg3d/code_models/worker_training.py

@@ -362,7 +362,9 @@ def log_parameters(self):
                 for k, v in d.items()
             ]
 
-    def train(self):
+    def train(
+        self, provided_model=None, provided_optimizer=None, provided_loss=None
+    ):
         try:
             if self.config is None:
                 self.config = config.WNetTrainingWorkerConfig()
@@ -395,11 +397,15 @@ def train(self):
             ###################################################
             self.log("- Getting the model")
             # Initialize the model
-            model = WNet(
-                in_channels=self.config.in_channels,
-                out_channels=self.config.out_channels,
-                num_classes=self.config.num_classes,
-                dropout=self.config.dropout,
+            model = (
+                WNet(
+                    in_channels=self.config.in_channels,
+                    out_channels=self.config.out_channels,
+                    num_classes=self.config.num_classes,
+                    dropout=self.config.dropout,
+                )
+                if provided_model is None
+                else provided_model
             )
             model.to(device)
 
@@ -458,7 +464,8 @@ def train(self):
                 optimizer = torch.optim.Adam(
                     model.parameters(), lr=self.config.learning_rate
                 )
-
+            if provided_optimizer is not None:
+                optimizer = provided_optimizer
             self.log("- Getting the loss functions")
             # Initialize the Ncuts loss function
             criterionE = SoftNCutsLoss(
@@ -538,6 +545,8 @@ def train(self):
                     beta = self.config.rec_loss_weight
 
                     loss = alpha * Ncuts + beta * reconstruction_loss
+                    if provided_loss is not None:
+                        loss = provided_loss
                     epoch_loss += loss.item()
                     # if WANDB_INSTALLED:
                     #     wandb.log({"Weighted sum of losses": loss.item()})