WIP autosave weights if worker has been canceled

Author: C-Achard

src/napari_cellseg3d/model_workers.py

@@ -768,6 +768,7 @@ def train(self):
                         "epoch": epoch,
                         "losses": epoch_loss_values,
                         "val_metrics": val_metric_values,
+                        "weights": model.state_dict(),
                     }
                     yield train_report
 
@@ -794,8 +795,8 @@ def train(self):
             f"at epoch: {best_metric_epoch}"
         )
         model.to("cpu")
-        optimizer = None
-        del optimizer
+        # optimizer = None
+        # del optimizer
         # del device
         # del model_id
         # del model_name

src/napari_cellseg3d/plugin_model_inference.py

@@ -1,5 +1,6 @@
 import os
 import warnings
+import numpy as np
 
 import napari
 # Qt
@@ -624,6 +625,9 @@ def on_yield(data, widget):
             )
 
             if data["instance_labels"] is not None:
+
+                widget.log.print_and_log(f"\nNUMBER OF CELLS : {np.amax(data['instance_labels'])}\n")
+
                 name = f"instance_labels_{image_id}"
                 instance_layer = viewer.add_labels(
                     data["instance_labels"], name=name

src/napari_cellseg3d/plugin_model_training.py

@@ -1,5 +1,6 @@
 import os
 import warnings
+import torch
 from pathlib import Path
 
 import matplotlib.pyplot as plt
@@ -168,6 +169,8 @@ def __init__(
         """Training worker for multithreading, should be a TrainingWorker instance from :doc:model_workers.py"""
         self.data = None
         """Data dictionary containing file paths"""
+        self.stop_requested = False
+        """Whether the worker should stop or not"""
 
         self.loss_dict = {
             "Dice loss": DiceLoss(sigmoid=True),
@@ -645,6 +648,10 @@ def start(self):
 
         """
 
+        if self.stop_requested:
+            self.log.print_and_log("Worker is already stopping !")
+            return
+
         if not self.check_ready():  # issues a warning if not ready
             err = "Aborting, please set all required paths"
             self.log.print_and_log(err)
@@ -734,6 +741,7 @@ def start(self):
             self.log.print_and_log(
                 f"Stop requested at {utils.get_time()}. \nWaiting for next validation step..."
             )
+            self.stop_requested = True
             self.btn_start.setText("Stopping... Please wait for next saving")
             self.worker.quit()
         else:
@@ -795,6 +803,10 @@ def on_yield(data, widget):
         )
         widget.update_loss_plot(data["losses"], data["val_metrics"])
 
+        if widget.stop_requested:
+            torch.save(data["weights"], os.path.join(widget.results_path, f"latest_weights_aborted_training_{utils.get_date_time()}.pth"))
+            widget.stop_requested = False
+
     # def clean_cache(self):
     #     """Attempts to clear memory after training"""
     #     # del self.worker