tried BCE loss + fixed saving errors in training + fixed incorrect plot clearing + tested Trailmap more

Author: C-Achard

src/napari_cellseg3d/model_workers.py

@@ -459,6 +459,7 @@ def log_parameters(self):
         self.log(f"Loss function is : {str(self.loss_function)}")
         self.log(f"Validation is performed every {self.val_interval} epochs")
         self.log(f"Batch size is {self.batch_size}")
+        self.log(f"Learning rate is {self.learning_rate}")
 
         if self.sampling:
             self.log(
@@ -661,14 +662,15 @@ def train(self):
         )
         print("\nDone")
 
+        print("Optimizer")
         optimizer = torch.optim.Adam(model.parameters(), self.learning_rate)
         dice_metric = DiceMetric(include_background=True, reduction="mean")
 
         best_metric = -1
         best_metric_epoch = -1
 
         # time = utils.get_date_time()
-
+        print("Weights")
         if self.weights_path is not None:
             if self.weights_path == "use_pretrained":
                 weights_file = model_class.get_weights_file()

src/napari_cellseg3d/plugin_model_training.py

@@ -1,6 +1,7 @@
 import os
 import warnings
 import torch
+from torch import nn
 from pathlib import Path
 
 import matplotlib.pyplot as plt
@@ -174,10 +175,12 @@ def __init__(
         """Data dictionary containing file paths"""
         self.stop_requested = False
         """Whether the worker should stop or not"""
+        self.start_time = ""
 
         self.loss_dict = {
             "Dice loss": DiceLoss(sigmoid=True),
             "Focal loss": FocalLoss(),
+            # "BCELoss":nn.BCELoss(),
             "Dice-Focal loss": DiceFocalLoss(sigmoid=True, lambda_dice=0.5),
             "Generalized Dice loss": GeneralizedDiceLoss(sigmoid=True),
             "DiceCELoss": DiceCELoss(sigmoid=True, lambda_ce=0.5),
@@ -226,6 +229,7 @@ def __init__(
         self.lbl_val_interv_choice = QLabel("Validation interval : ", self)
 
         self.learning_rate_dict = {
+            "1e-2": 1e-2,
             "1e-3": 1e-3,
             "1e-4": 1e-4,
             "1e-5": 1e-5,
@@ -238,6 +242,7 @@ def __init__(
         ) = ui.make_combobox(
             self.learning_rate_dict.keys(), label="Learning rate"
         )
+        self.learning_rate_choice.setCurrentIndex(1)
 
         self.augment_choice = ui.make_checkbox("Augment data")
 
@@ -676,7 +681,7 @@ def start(self):
         Returns: Returns empty immediately if the file paths are not set correctly.
 
         """
-
+        self.start_time = utils.get_time_filepath()
         if self.stop_requested:
             self.log.print_and_log("Worker is already stopping !")
             return
@@ -720,8 +725,11 @@ def start(self):
 
             self.results_path = (
                 self.results_path
-                + f"/{model_dict['name']}_results_{utils.get_date_time()}"
+                + f"/{model_dict['name']}_results_{self.start_time}"
             )
+            os.makedirs(
+                self.results_path, exist_ok=False
+            )  # avoid overwrite where possible
 
             if self.use_transfer_choice.isChecked():
                 if self.custom_weights_choice.isChecked():
@@ -731,10 +739,6 @@ def start(self):
             else:
                 weights_path = None
 
-            os.makedirs(
-                self.results_path, exist_ok=False
-            )  # avoid overwrite where possible
-
             self.log.print_and_log(
                 f"Notice : Saving results to : {self.results_path}"
             )
@@ -782,13 +786,11 @@ def start(self):
 
     def on_start(self):
         """Catches started signal from worker"""
-        if self.plot_dock is not None:
-            self._viewer.window.remove_dock_widget(self.plot_dock)
-            self.plot_dock = None
 
+        self.remove_docked_widgets()
         self.display_status_report()
 
-        self.log.print_and_log(f"Worker started at {utils.get_time()}")
+        self.log.print_and_log(f"Worker started at {self.start_time}")
         self.log.print_and_log("\nWorker is running...")
 
     def on_finish(self):
@@ -801,7 +803,7 @@ def on_finish(self):
             self.canvas.figure.savefig(
                 (
                     self.results_path
-                    + f"/final_metric_plots_{utils.get_date_time()}.png"
+                    + f"/final_metric_plots_{utils.get_time_filepath()}.png"
                 ),
                 format="png",
             )
@@ -817,11 +819,14 @@ def on_finish(self):
 
         self.worker = None
         self.empty_cuda_cache()
+
+        self.results_path = ""
         # self.clean_cache() # trying to fix memory leak
 
     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.empty_cuda_cache()
         # self.clean_cache()
 
@@ -840,7 +845,7 @@ def on_yield(data, widget):
                 data["weights"],
                 os.path.join(
                     widget.results_path,
-                    f"latest_weights_aborted_training_{utils.get_date_time()}.pth",
+                    f"latest_weights_aborted_training_{utils.get_time()}.pth",
                 ),
             )
             widget.stop_requested = False