added deterministic

Author: C-Achard

src/napari_cellseg3d/model_workers.py

@@ -29,6 +29,9 @@
 from monai.transforms import RandSpatialCropSamplesd
 from monai.transforms import SpatialPadd
 from monai.transforms import Zoom
+from monai.utils import set_determinism
+
+# threads
 from napari.qt.threading import GeneratorWorker
 from napari.qt.threading import WorkerBaseSignals
 
@@ -388,6 +391,7 @@ def __init__(
         num_samples,
         sample_size,
         do_augmentation,
+        deterministic,
     ):
         """Initializes a worker for inference with the arguments needed by the :py:func:`~train` function.
 
@@ -420,6 +424,8 @@ def __init__(
 
             * do_augmentation : whether to perform data augmentation or not
 
+            * deterministic : dict with "use deterministic" : bool, whether to use deterministic training, "seed": seed for RNG
+
            Note: See :py:func:`~train`
         """
 
@@ -443,6 +449,7 @@ def __init__(
         self.sample_size = sample_size
 
         self.do_augment = do_augmentation
+        self.seed_dict = deterministic
 
     def log(self, text):
         """Sends a signal that ``text`` should be logged
@@ -454,7 +461,12 @@ def log(self, text):
 
     def log_parameters(self):
 
-        self.log("\nParameters summary :")
+        self.log("\nParameters summary :\n")
+
+        if self.seed_dict["use deterministic"]:
+            self.log(f"Deterministic training is enabled")
+            self.log(f"Seed is {self.seed_dict['seed']}")
+
         self.log(f"Training for {self.max_epochs} epochs")
         self.log(f"Loss function is : {str(self.loss_function)}")
         self.log(f"Validation is performed every {self.val_interval} epochs")
@@ -474,6 +486,8 @@ def log_parameters(self):
         if self.weights_path is not None:
             self.log(f"Using weights from : {self.weights_path}")
 
+        self.log("\n")
+
     def train(self):
         """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.
@@ -508,13 +522,20 @@ def train(self):
         * sample_size : the size of the patches to extract when sampling
 
         * do_augmentation : whether to perform data augmentation or not
+
+        * deterministic : dict with "use deterministic" : bool, whether to use deterministic training, "seed": seed for RNG
         """
 
         #########################
         # error_log = open(results_path +"/error_log.log" % multiprocessing.current_process().name, 'x')
         # faulthandler.enable(file=error_log, all_threads=True)
         #########################
 
+        if self.seed_dict["use deterministic"]:
+            set_determinism(
+                seed=self.seed_dict["seed"]
+            )  # use_deterministic_algorithms = True causes cuda error
+
         sys = platform.system()
         print(sys)
         if sys == "Darwin":  # required for macOS ?

src/napari_cellseg3d/plugin_model_training.py

@@ -273,6 +273,15 @@ def __init__(
             "Transfer weights", self.toggle_transfer_param
         )
 
+        self.use_deterministic_choice = ui.make_checkbox(
+            "Deterministic training", func=self.toggle_deterministic_param
+        )
+        self.box_seed = ui.make_n_spinboxes(max=10000000, default=23498)
+        self.lbl_seed = ui.make_label("Seed", self)
+        self.container_seed = ui.combine_blocks(
+            self.box_seed, self.lbl_seed, horizontal=False
+        )
+
         self.progress = QProgressBar()
         self.progress.setVisible(False)
         """Dock widget containing the progress bar"""
@@ -304,6 +313,12 @@ def toggle_transfer_param(self):
         else:
             self.custom_weights_choice.setVisible(False)
 
+    def toggle_deterministic_param(self):
+        if self.use_deterministic_choice.isChecked():
+            self.container_seed.setVisible(True)
+        else:
+            self.container_seed.setVisible(False)
+
     def check_ready(self):
         """
         Checks that the paths to the images and labels are correctly set
@@ -595,6 +610,22 @@ def build(self):
         train_tab_layout.addWidget(train_param_group_w)
         # end of training params group
         ##################
+        ui.add_blank(self, train_tab_layout)
+        ##################
+        # deterministic choice group
+        seed_w, seed_l = ui.make_group(
+            "Deterministic training", R=1, B=5, T=11
+        )
+
+        seed_l.addWidget(self.use_deterministic_choice, alignment=ui.LEFT_AL)
+        seed_l.addWidget(self.container_seed, alignment=ui.LEFT_AL)
+        self.container_seed.setVisible(False)
+
+        seed_w.setLayout(seed_l)
+        train_tab_layout.addWidget(seed_w)
+
+        # end of deterministic choice group
+        ##################
         # buttons
 
         ui.add_blank(self, train_tab_layout)
@@ -710,7 +741,15 @@ def start(self):
             self.data = self.create_train_dataset_dict()
             self.max_epochs = self.epoch_choice.value()
 
-            self.learning_rate = self.learning_rate_dict[self.learning_rate_choice.currentText()]
+            self.learning_rate = self.learning_rate_dict[
+                self.learning_rate_choice.currentText()
+            ]
+
+            seed_dict = {
+                "use deterministic": self.use_deterministic_choice.isChecked(),
+                "seed": self.box_seed.value(),
+            }
+
 
             self.patch_size = []
             [
@@ -722,10 +761,9 @@ def start(self):
                 "class": self.get_model(self.model_choice.currentText()),
                 "name": self.model_choice.currentText(),
             }
-
             self.results_path = (
-                self.results_path
-                + f"/{model_dict['name']}_results_{self.start_time}"
+                    self.results_path
+                    + f"/{model_dict['name']}_results_{utils.get_date_time()}"
             )
             os.makedirs(
                 self.results_path, exist_ok=False
@@ -758,6 +796,7 @@ def start(self):
                 num_samples=self.num_samples,
                 sample_size=self.patch_size,
                 do_augmentation=self.augment_choice.isChecked(),
+                deterministic=seed_dict,
             )
 
             [btn.setVisible(False) for btn in self.close_buttons]
@@ -826,7 +865,7 @@ def on_finish(self):
     def on_error(self):
         """Catches errored signal from worker"""
         self.log.print_and_log(f"WORKER ERRORED at {utils.get_time()}")
-        self.worker=None
+        self.worker = None
         self.empty_cuda_cache()
         # self.clean_cache()