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

CHANGELOG.md

@@ -5,6 +5,21 @@ All notable chagnes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),  
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.9] - 2026-09-11
+
+### Added
+
+
+#### Initial training resume support (`virtual_stain_flow/vsf_logging/`, `virtual_stain_flow/trainers/`)
+- Added MLflow logging for optimizer state to support resuming training runs from logged artifacts and trainer state.
+
+### Fixed
+
+#### Pandas validation API update (`virtual_stain_flow/datasets/ds_engine/`)
+- Replaced deprecated `DataFrame.applymap` usage with `DataFrame.map` in dataset input validation to keep the GH Actions test suite green.
+
+
+
 ---
 
 ## [0.4.8] - 2026-08-03

pyproject.toml

@@ -6,13 +6,13 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "virtual_stain_flow"
-version = "0.4.7"
+version = "0.4.9"
 description = "For developing virtual staining models"
 requires-python = ">=3.9"
 dependencies = [
   "pyarrow",
   "numpy",
-  "pandas",
+  "pandas>=2.1",
   "seaborn",
   "matplotlib",
   "scikit-learn",

src/virtual_stain_flow/datasets/ds_engine/input_validation.py

@@ -16,7 +16,7 @@
 def _cell_contains_pathlike(x: Any, *, check_exists: bool) -> bool:
     """
     Validate that a cell contains a valid path-like object.
-    Intended to be used as a lambda function in the DataFrame.applymap method.
+    Intended to be used as a lambda function in the DataFrame.map method.
 
     :param x: The cell value to check.
     :param check_exists: Whether to check if the path exists.

src/virtual_stain_flow/datasets/ds_engine/manifest.py

@@ -179,12 +179,12 @@ def get_image_dimensions(
 
     def _serialize_file_index(self) -> pd.DataFrame:
         """Serialize file_index to pd.DataFrame"""
-        return self.file_index.copy().applymap(lambda x: str(x))
+        return self.file_index.map(str)
 
     @staticmethod
     def _deserialize_file_index(file_index: pd.DataFrame) -> pd.DataFrame:
         """Deserialize file_index from pd.DataFrame"""
-        return file_index.applymap(lambda x: Path(x))
+        return file_index.map(Path)
 
     def to_config(self) -> Dict[str, Any]:
         """Serialize to dict"""

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: int = 0,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
@@ -68,6 +69,8 @@ def __init__(
             dataset is provided. Default is 0.15.
         :param metrics: Dictionary of metrics to be logged.
         :param device: (optional) The device to be used for training.
+        :param epoch: (optional) The starting epoch for training. 
+            Useful for resuming training from a checkpoint.
         :param early_termination_metric: (optional) The metric to update 
             early-termination count on the validation dataset. 
             If None, early termination is disabled and the
@@ -102,20 +105,25 @@ def __init__(
             **kwargs
         )
         self._init_state(
+            epoch,
             early_termination_metric, early_termination_mode, **kwargs)
 
     def _init_state(
         self, 
+        epoch: int,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs
     ):
 
+        if epoch is None:
+            raise TypeError("epoch must be an integer, not None.")
+
         # Epoch state
-        self._epoch = 0
+        self._epoch = epoch
 
         # Progress tracking for loss weight scheduling
-        self._progress = Progress(epoch=0, step=0)
+        self._progress = Progress(epoch=epoch, step=0)
 
         # Loss and metrics state
         self._train_losses = defaultdict(list)
@@ -183,7 +191,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,8 +201,24 @@ 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 = self._get_dataset_size(self._train_loader)
+        self._val_n = self._get_dataset_size(self._val_loader)
+        self._test_n = self._get_dataset_size(self._test_loader)
+
         return None
 
+    @staticmethod
+    def _get_dataset_size(loader) -> Optional[int]:
+        dataset = getattr(loader, 'dataset', None)
+        if dataset is None:
+            return None
+
+        try:
+            return len(dataset)
+        except TypeError:
+            return None
+
     @abstractmethod
     def train_step(self, inputs: torch.Tensor, targets: torch.Tensor)->Dict[str, float]:
         """
@@ -323,17 +346,17 @@ def train(
         if hasattr(logger, "on_train_start"):
             logger.on_train_start()
 
-        self._epochs = epochs
+        epoch_range = range(self.epoch + 1, self.epoch + epochs + 1)
         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)
+            epoch_range, desc="Training", unit="epoch") if verbose else None
+        iterable = self._epoch_pbar if self._epoch_pbar else epoch_range
 
         self._early_stop_helper.initialize_early_stop(patience=patience if patience else epochs)
 
         for epoch in iterable:
 
-            # Increment the epoch counter
-            self.epoch += 1
+            # Synchronize trainer state for loggers, callbacks, and schedulers
+            self.epoch = epoch
 
             # Invoke the on_epoch_start method of the logge
             if hasattr(logger, "on_epoch_start"):
@@ -406,20 +429,44 @@ def _update_epoch_progress(
     def save_model(
         self,
         save_path: pathlib.Path,
-        file_name_prefix: Optional[str] = None,
+        file_name_prefix: str = 'generator',
         file_name_suffix: Optional[str] = None,
         file_ext: str = '.pth',
         best_model: bool = True
     ) -> Optional[List[pathlib.Path]]:
         return save_model(
             self,
             save_path=save_path,
-            file_name_prefix=file_name_prefix or 'generator',
+            file_name_prefix=file_name_prefix,
             file_name_suffix=file_name_suffix,
             file_ext=file_ext,
             save_best_model=best_model
         )
 
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: str = 'optimizer',
+        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 +499,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 +527,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

@@ -19,7 +19,6 @@ class TrainerProtocol(Protocol):
     """
 
     _batch_size: int
-    _epochs: int
     _patience: int
     _device: torch.device
 
@@ -44,6 +43,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 +70,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: str = 'generator',
+        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: str = 'optimizer',
+        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_optimizer.py

@@ -0,0 +1,36 @@
+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: str = 'optimizer',
+    file_name_suffix: Optional[str] = None,
+    file_ext: str = '.pth',
+) -> List[Path]:
+
+    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 []
+