Add auto logging of trainer metadata and optimizer state saving to help with resumation of training

src/virtual_stain_flow/trainers/AbstractTrainer.py

@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader
 
 from .trainer_protocol import TrainerProtocol
-from .trainer_utils import EarlyStopHelper, save_model
+from .trainer_utils import EarlyStopHelper, save_model, save_optimizer_state
 from ..metrics.AbstractMetrics import AbstractMetrics
 from ..engine.progress import Progress
 from ..datasets.data_split import default_random_split
@@ -42,6 +42,7 @@ def __init__(
         test_ratio: Optional[float] = 0.15,
         metrics: Dict[str, AbstractMetrics] = None,
         device: Optional[torch.device] = None,
+        epoch: Optional[int] = 0,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
@@ -102,17 +103,19 @@ def __init__(
             **kwargs
         )
         self._init_state(
+            epoch,
             early_termination_metric, early_termination_mode, **kwargs)
 
     def _init_state(
         self, 
+        epoch,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs
     ):
 
         # Epoch state
-        self._epoch = 0
+        self._epoch = epoch
 
         # Progress tracking for loss weight scheduling
         self._progress = Progress(epoch=0, step=0)
@@ -183,7 +186,6 @@ def _init_data(
                 **kwargs
             )
 
-            self._batch_size = batch_size
             self._train_ratio, self._val_ratio, self._test_ratio = (
                 train_ratio, val_ratio, test_ratio
             )
@@ -194,6 +196,11 @@ def _init_data(
                 "or provide at least train_loader."
             )
 
+        self._batch_size = self._train_loader.batch_size if hasattr(self._train_loader, 'batch_size') else None
+        self._train_n = len(self._train_loader.dataset) if hasattr(self._train_loader, 'dataset') else None
+        self._val_n = len(self._val_loader.dataset) if hasattr(self._val_loader, 'dataset') else None
+        self._test_n = len(self._test_loader.dataset) if hasattr(self._test_loader, 'dataset') else None
+
         return None
 
     @abstractmethod
@@ -323,7 +330,6 @@ def train(
         if hasattr(logger, "on_train_start"):
             logger.on_train_start()
 
-        self._epochs = epochs
         self._epoch_pbar: Optional[tqdm] = tqdm(
             range(epochs), desc="Training", unit="epoch") if verbose else None
         iterable = self._epoch_pbar if self._epoch_pbar else range(epochs)
@@ -420,6 +426,30 @@ def save_model(
             save_best_model=best_model
         )
 
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True
+    ) -> Optional[List[pathlib.Path]]:
+        """
+        Save the optimizer state to the specified path.
+        """
+        if not recent:
+            raise NotImplementedError(
+                "Saving non-recent optimizer states is not implemented yet."
+            )
+        file_name_suffix = file_name_suffix or 'recent'
+        return save_optimizer_state(
+            trainer=self,
+            save_path=save_path,
+            file_name_prefix=file_name_prefix,
+            file_name_suffix=file_name_suffix,
+            file_ext=file_ext
+        )
+
     """
     Log property
     """
@@ -452,6 +482,18 @@ def val_ratio(self):
     @property
     def test_ratio(self):
         return self._test_ratio
+
+    @property
+    def train_n(self):
+        return self._train_n
+
+    @property
+    def val_n(self):
+        return self._val_n
+
+    @property
+    def test_n(self):
+        return self._test_n
 
     @property
     def model(self):
@@ -468,11 +510,7 @@ def device(self):
     @property
     def batch_size(self):
         return self._batch_size
-
-    @property
-    def epochs(self):
-        return self._epochs
-
+
     @property
     def patience(self):
         return self._patience

src/virtual_stain_flow/trainers/logging_gan_trainer.py

@@ -84,7 +84,6 @@ def __init__(
         super().__init__(
             model=generator, # register generator as main model for early stopping
             optimizer=generator_optimizer,
-            losses=generator_loss_group,
             device=device,
             **kwargs
         )

src/virtual_stain_flow/trainers/trainer_protocol.py

@@ -44,6 +44,18 @@ def train(self, *args: Any, **kwargs: Any) -> None: ...
     @property
     def epoch(self) -> int: ...
 
+    @property
+    def batch_size(self) -> int: ...
+
+    @property
+    def train_n(self) -> int: ...
+
+    @property
+    def val_n(self) -> int: ...
+
+    @property
+    def test_n(self) -> int: ...
+
     @property
     def device(self) -> torch.device: ...
 
@@ -59,9 +71,19 @@ def best_model(self) -> Optional[torch.nn.Module]: ...
     def save_model(
         self, 
         save_path: pathlib.Path, 
-        file_name_prefix: Optional[str], 
-        file_name_suffix: Optional[str], 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
         file_ext: str = '.pth',
         best_model: bool = True,
     ) -> Optional[List[pathlib.Path]]:
         ...
+
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True,
+    ) -> Optional[List[pathlib.Path]]:
+        ...

src/virtual_stain_flow/trainers/trainer_utils/__init__.py

@@ -7,10 +7,12 @@
     _get_latest_metric_value,
 )
 from .save_model import save_model
+from .save_optimizer import save_optimizer_state
 
 
 __all__ = [
     "EarlyStopHelper",
     "_get_latest_metric_value",
     "save_model",
+    "save_optimizer_state",
 ]

src/virtual_stain_flow/trainers/trainer_utils/save_model.py

@@ -6,12 +6,15 @@
 def save_model(
     trainer: 'TrainerProtocol',
     save_path: Path,
-    file_name_prefix: str = 'generator',
+    file_name_prefix: Optional[str] = None,
     file_name_suffix: Optional[str] = None,
     file_ext: str = '.pth',
     save_best_model: bool = True,
 ) -> List[Path]:
 
+    if file_name_prefix is None:
+        file_name_prefix = 'generator'
+
     if file_name_suffix is None:
         file_name_suffix = 'weights_' + (
             'best' if save_best_model else str(trainer.epoch)

src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import Optional, List
+
+import torch
+
+from ..trainer_protocol import TrainerProtocol
+
+
+def save_optimizer_state(
+    trainer: 'TrainerProtocol',
+    save_path: Path,
+    file_name_prefix: Optional[str] = None,
+    file_name_suffix: Optional[str] = None,
+    file_ext: str = '.pth',
+) -> List[Path]:
+
+    if file_name_prefix is None:
+        file_name_prefix = 'optimizer'
+
+    if file_name_suffix is None:
+        file_name_suffix = f"{trainer.epoch}"
+
+    optimizer = trainer.optimizer
+
+    if optimizer is None:
+        return []
+
+    save_file = save_path / f"{file_name_prefix}_{file_name_suffix}{file_ext}"
+
+    torch.save(
+        optimizer.state_dict(), 
+        save_file
+    )
+
+    if save_file.exists():
+        return [save_file]
+
+    return []
+