From b8000e8e44cb9d31f72891bd8bbf89b90e4d104e Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 11:56:08 -0600
Subject: [PATCH 01/10] Add optimizer state saving functionality as new trainer
 contract, update existing properties by best effort inferring some. Remove
 some confusing and unused properties. Adjust model saving defaults.

---
 .../trainers/AbstractTrainer.py               | 56 ++++++++++++++++---
 .../trainers/logging_gan_trainer.py           |  1 -
 .../trainers/trainer_protocol.py              | 26 ++++++++-
 .../trainers/trainer_utils/__init__.py        |  2 +
 .../trainers/trainer_utils/save_model.py      |  5 +-
 .../trainers/trainer_utils/save_optimizer.py  | 39 +++++++++++++
 tests/trainers/test_abstract_trainer.py       |  5 +-
 7 files changed, 119 insertions(+), 15 deletions(-)
 create mode 100644 src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 932faf3..82bf209 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/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
diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
index e51cc14..ca9a25e 100644
--- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py
+++ b/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
         )
diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index 5fc3cb0..e0de924 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/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]]:
+        ...
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
index f0a09fa..a706015 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
+++ b/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",
 ]
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
index 528d81b..8063199 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
+++ b/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)
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
new file mode 100644
index 0000000..97e899c
--- /dev/null
+++ b/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 []
+ 
diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py
index 2cb9099..0d22901 100644
--- a/tests/trainers/test_abstract_trainer.py
+++ b/tests/trainers/test_abstract_trainer.py
@@ -697,8 +697,9 @@ def test_batch_size_property_with_loader_init(self, minimal_model, minimal_optim
             device=torch.device('cpu')
         )
         
-        # When providing loaders, batch_size is set to None
-        assert trainer.batch_size is None
+        # When providing loaders, batch_size is inferred from the train loader
+        # should therefore match
+        assert trainer.batch_size is train_dataloader.batch_size
     
     def test_batch_size_property_default_value(self, minimal_model, minimal_optimizer, dataset_for_splitting):
         """Verify that batch_size property uses default value when not specified."""

From 3c199222918c29d3ff4e5de9d8d6be8161dcc2fa Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 12:01:52 -0600
Subject: [PATCH 02/10] Implement AutoTrainerLogger for adding auto logging of
 certain useful trainer metadata and couple logging of model weights by
 logging of optimizer state to facilitate resuming of complete trainings in
 whole epoch increments.

---
 .../vsf_logging/MlflowLogger.py               | 123 +++---------------
 .../vsf_logging/auto_loggers/__init__.py      |   2 +
 .../auto_loggers/optimizer_config_logger.py   |  12 +-
 .../auto_loggers/trainer_config_logger.py     |  35 +++++
 .../vsf_logging/logger_utils/__init__.py      |   6 +
 .../vsf_logging/logger_utils/log_artifacts.py |  83 ++++++++++++
 6 files changed, 155 insertions(+), 106 deletions(-)
 create mode 100644 src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
 create mode 100644 src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
 create mode 100644 src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py

diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
index 8933283..c534378 100644
--- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
+++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
@@ -15,7 +15,9 @@
     AutoLossGroupConfigLogger,
     AutoModelConfigLogger,
     AutoOptimizerConfigLogger,
+    AutoTrainerLogger
 )
+from .logger_utils import _log_artifact, _log_trainer_artifact
 from .callbacks.LoggerCallback import (
     AbstractLoggerCallback,
     log_type
@@ -145,6 +147,7 @@ def __init__(
         self._model_config_logger = AutoModelConfigLogger(self)
         self._optimizer_config_logger = AutoOptimizerConfigLogger(self)
         self._loss_group_config_logger = AutoLossGroupConfigLogger(self)
+        self._trainer_logger = AutoTrainerLogger(self)
 
         return None
     
@@ -208,7 +211,7 @@ def on_train_start(self):
         self._model_config_logger.log_model_configs(self.trainer)
         self._optimizer_config_logger.log_optimizer_configs(self.trainer)
         self._loss_group_config_logger.log_loss_group_configs(self.trainer)
-
+        self._trainer_logger.log_trainer_config(self.trainer)
         
         for callback in self.callbacks:
             # TODO consider if we want hasattr checks
@@ -248,16 +251,10 @@ def on_epoch_end(self):
 
         if self._save_model_every_n_epochs is not None:
             if self.trainer.epoch % self._save_model_every_n_epochs == 0:
-                self._save_model_weights(
-                    artifact_path='weights',
-                    best_model=False
-                )
+                _log_trainer_artifact(self.trainer, best_model=False)
 
         if self._save_best_model:
-            self._save_model_weights(
-                artifact_path='weights',
-                best_model=True
-            )
+            _log_trainer_artifact(self.trainer, best_model=True)
 
         # Call on_epoch_end for all registered callbacks
         for callback in self.callbacks:
@@ -323,14 +320,14 @@ def end_run(self):
             print("No active MLflow run to end.")
 
     """
-    Exposed? logging methods
+    Exposed logging methods
     """    
     def log_artifact(
-            self,
-            tag: str,
-            file_path: pathlib.Path,
-            stage: Optional[str] = None
-        ):
+        self,
+        tag: str,
+        file_path: pathlib.Path,
+        stage: Optional[str] = None
+    ):
         """
         Log an artifact to MLflow.
 
@@ -339,29 +336,11 @@ def log_artifact(
         :param stage: Optional stage to categorize the artifact, defaults to None.
         :raises TypeError: If file_path is not a pathlib.Path instance.
         """
-
-        if not isinstance(file_path, pathlib.Path):
-            raise TypeError("file_path must be a pathlib.Path instance.")
-
-        artifact_path = ''
-        artifact_ext = file_path.suffix.lower()
-        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
-            # log as plot artifact
-            artifact_path += 'plots/'
-        elif artifact_ext in ['.pth', '.pt']:
-            # log as model artifact
-            artifact_path += 'weights/'
-        else:
-            # log as generic artifact
-            artifact_path += 'artifacts/'
-        
-        if stage is not None:
-            artifact_path += f"{stage}/"
-        artifact_path += f"{tag}"
-
-        mlflow.log_artifact(
-            str(file_path), 
-            artifact_path=artifact_path
+        log_subdirs = [stage] if stage is not None else []
+        log_subdirs = log_subdirs + [tag] if tag is not None else log_subdirs
+        _log_artifact(
+            file_path=file_path,
+            artifact_subdirs=log_subdirs
         )
 
     def log_metric(
@@ -464,31 +443,6 @@ def _log_callback_output(
                 continue
                 # raise TypeError("Unsupported callback return type for logging.")
 
-    def _save_model_weights(
-        self,
-        prefix: Optional[str] = None,
-        suffix: Optional[str] = None,
-        artifact_path: str = "weights",
-        best_model: bool = True
-    ):
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            
-            tmpdirpath = pathlib.Path(tmpdirname)
-            
-            saved_file_paths = self.trainer.save_model(
-                save_path=tmpdirpath,
-                file_name_prefix=prefix,
-                file_name_suffix=suffix,
-                file_ext='.pth',
-                best_model=best_model
-            )
-
-            for saved_file_path in (saved_file_paths or []):
-                mlflow.log_artifact(
-                    str(saved_file_path), 
-                    artifact_path=artifact_path
-                )
-
     def log_config(
         self,
         tag: str,
@@ -550,7 +504,6 @@ def __check_trainer_bound(
         if self.trainer is None:
             raise RuntimeError("No trainer bound to logger. Cannot access trainer attributes.")
 
-
     def get_epoch(
             self
         ) -> int:
@@ -595,48 +548,8 @@ def get_model(
     @property
     def run_id(self):
         return self._run_id
-    
-    """
-    Unimplemented helper that might be useful
-    """
-    def _log_dict_as_yaml(
-        self,
-        dict: Dict[str, Any],
-    ):
-        """
-        Log a dictionary as a YAML file in MLflow.
-        
-        :param dict: The dictionary to log.
-        """
-        
-        # TODO implement this method to convert dict to YAML and log it
-        raise NotImplementedError(
-            "log_dict_as_yaml method is not implemented. "
-            "Please implement this method to log dictionary as YAML."
-        )
-    
-    def _log_dict_as_param(
-        self
-    ):
-        """
-        Log a dictionary as parameters in MLflow.
-        
-        :return: None
-        :raises NotImplementedError: If the method is not implemented.
-        
-        This method is intended to log a dictionary as parameters in MLflow.
-        It is currently not implemented and raises a NotImplementedError.
-        """
-
-        # 1. DO flatten dict
-
-        # 2. DO mlflow log
-        
-        raise NotImplementedError(
-            "log_dict_as_param method is not implemented. "
-            "Please implement this method to log dictionary as parameters."
-        )
 
+    
     """
     Overridden destructor method to ensure MLflow run is ended
     """
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
index 2860937..2625d41 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
@@ -1,9 +1,11 @@
 from .loss_group_config_logger import AutoLossGroupConfigLogger
 from .model_config_logger import AutoModelConfigLogger
 from .optimizer_config_logger import AutoOptimizerConfigLogger
+from .trainer_config_logger import AutoTrainerLogger
 
 __all__ = [
     "AutoModelConfigLogger",
     "AutoOptimizerConfigLogger",
     "AutoLossGroupConfigLogger",
+    "AutoTrainerLogger",
 ]
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
index fa3f344..c95537e 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
@@ -1,4 +1,5 @@
 from typing import Any, Dict, List, Optional
+import inspect
 
 import mlflow
 from torch.optim import Optimizer
@@ -44,13 +45,22 @@ def log_optimizer_configs(
                 continue
 
             try:
+                
+                defaults = dict(optimizer.defaults)
+                init_signature = inspect.signature(optimizer.__class__.__init__)
+                valid_params = init_signature.parameters.keys()
+                
                 opt_config: Optional[Dict[str, Any]] = {
                     "class_path": (
                         f"{optimizer.__class__.__module__}."
                         f"{optimizer.__class__.__name__}"
                     ),
-                    "defaults": dict(optimizer.defaults),
+                    "defaults": defaults,
+                    "init": {
+                        k: v for k, v in defaults.items() if k in valid_params
+                    }
                 }
+
             except Exception as e:
                 print(f"Could not get optimizer config for logging: {e}")
                 opt_config = None
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
new file mode 100644
index 0000000..cf6c292
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
@@ -0,0 +1,35 @@
+from typing import Any, Optional
+
+from ...trainers.trainer_protocol import TrainerProtocol
+
+
+class AutoTrainerLogger:
+    """
+    Auto-log trainer metadata to MLflow.
+    """
+
+    def __init__(self, logger: Any) -> None:
+        self._logger = logger
+
+    def log_trainer_config(self, trainer: Optional[TrainerProtocol]) -> None:
+        
+        if trainer is None:
+            return
+
+        config = {
+            "class_path": f"{trainer.__class__.__module__}.{trainer.__class__.__name__}",
+            "device": str(trainer.device), # device used for training
+            "batch_size": trainer.batch_size, # batch size used for training
+            "train_n": trainer.train_n,
+            "val_n": trainer.val_n,
+            "test_n": trainer.test_n,
+        }
+
+        try:
+            self._logger.log_config(
+                tag="trainer",
+                config=config,
+                stage=None,
+            )
+        except Exception as e:
+            print(f"Could not log trainer config: {e}")
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
new file mode 100644
index 0000000..04ad9c8
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
@@ -0,0 +1,6 @@
+from .log_artifacts import _log_artifact, _log_trainer_artifact
+
+__all__ = [
+    "_log_artifact",
+    "_log_trainer_artifact",
+]
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
new file mode 100644
index 0000000..199fe7c
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
@@ -0,0 +1,83 @@
+from typing import Optional, List
+import pathlib
+import tempfile
+
+import mlflow
+
+from virtual_stain_flow.trainers.trainer_protocol import TrainerProtocol
+
+
+def _log_artifact(
+    file_path: pathlib.Path,
+    artifact_path: Optional[str] = None,
+    artifact_subdirs: Optional[List[str]] = None
+) -> None:
+    """
+    Logs a single artifact to MLflow.
+
+    :param file_path: The path to the file to log as an artifact.
+    :param artifact_path: Optional artifact path within the MLflow run, defaults to None.
+    :param artifact_subdirs: Optional list of subdirectories to include in the artifact path, defaults to None.
+    :raises TypeError: If file_path is not a pathlib.Path instance.
+    """
+
+    if not isinstance(file_path, pathlib.Path):
+        raise TypeError("file_path must be a pathlib.Path instance.")
+
+    if artifact_path is None:
+        artifact_ext = file_path.suffix.lower()
+        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
+            # log as plot artifact
+            artifact_path = 'plots'
+        elif artifact_ext in ['.pth', '.pt']:
+            # log as model artifact
+            artifact_path = 'weights'
+        else:
+            # log as generic artifact
+            artifact_path = 'artifacts'
+
+    artifact_subdirs = [] if artifact_subdirs is None else artifact_subdirs
+
+    path_parts = [artifact_path, *artifact_subdirs]
+    clean_parts = [
+        part.replace('\\', '/').strip('/')
+        for part in path_parts
+        if part and part.strip('/\\')
+    ]
+
+    # Build a lexical POSIX path for MLflow without resolving it on the host OS.
+    artifact_path = str(pathlib.PurePosixPath(*clean_parts)) if clean_parts else ''
+
+    mlflow.log_artifact(str(file_path), artifact_path=artifact_path)
+
+
+def _log_trainer_artifact(
+    trainer: "TrainerProtocol",
+    best_model: bool = True,
+) -> None:
+    """
+    Logs the model and optimizer state artifacts from a trainer to MLflow.
+    The most recent optimizer state is saved and logged regardless of the best_model flag.
+
+    :param trainer: The trainer instance adhering to TrainerProtocol.
+    :param best_model: Whether to log only the best model, defaults to True.
+    :raises TypeError: If the provided trainer does not adhere to TrainerProtocol.
+    """
+
+    if not isinstance(trainer, TrainerProtocol):
+        raise TypeError("The provided trainer must adhere to the TrainerProtocol.")
+    
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        
+        tmpdirpath = pathlib.Path(tmpdirname)
+        saved_model_paths = trainer.save_model(
+            save_path=tmpdirpath, best_model=best_model
+        )
+        for saved_model_path in (saved_model_paths or []):
+            _log_artifact(saved_model_path, artifact_path='weights')
+
+        saved_optimizer_paths = trainer.save_optimizer_state(
+            save_path=tmpdirpath, recent=True
+        )
+        for saved_optimizer_path in (saved_optimizer_paths or []):
+            _log_artifact(saved_optimizer_path, artifact_path='optimizer')

From bd3036b49eabf703a4c0a06ecf9a06d7e4839c7c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 6 Aug 2026 12:12:53 -0600
Subject: [PATCH 03/10] Add unit tests for model and optimizer state saving
 functionality

---
 .../trainers/trainer_utils/test_save_model.py | 23 +++++++++++++++++++
 .../trainer_utils/test_save_optimizer.py      | 23 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 tests/trainers/trainer_utils/test_save_model.py
 create mode 100644 tests/trainers/trainer_utils/test_save_optimizer.py

diff --git a/tests/trainers/trainer_utils/test_save_model.py b/tests/trainers/trainer_utils/test_save_model.py
new file mode 100644
index 0000000..be65008
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_model.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_model helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_model import save_model
+
+
+def test_save_model_returns_empty_when_target_model_is_none(tmp_path):
+    trainer = SimpleNamespace(model=None, best_model=None, epoch=2)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=True)
+
+    assert saved_paths == []
+
+
+def test_save_model_saves_current_model_with_default_name(mock_model_with_save, tmp_path):
+    trainer = SimpleNamespace(model=mock_model_with_save, best_model=None, epoch=3)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=False)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "generator_weights_3.pth"
diff --git a/tests/trainers/trainer_utils/test_save_optimizer.py b/tests/trainers/trainer_utils/test_save_optimizer.py
new file mode 100644
index 0000000..dac7d52
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_optimizer.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_optimizer_state helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_optimizer import save_optimizer_state
+
+
+def test_save_optimizer_state_returns_empty_for_missing_optimizer(tmp_path):
+    trainer = SimpleNamespace(optimizer=None, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert saved_paths == []
+
+
+def test_save_optimizer_state_saves_with_default_name(minimal_optimizer, tmp_path):
+    trainer = SimpleNamespace(optimizer=minimal_optimizer, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "optimizer_7.pth"

From 768907113c62eec9e14ad3aabff0bbd4ebb8ffa6 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Fri, 7 Aug 2026 17:44:05 -0600
Subject: [PATCH 04/10] Fix MlflowLogger to use _log_trainer_artifact for model
 weight logging and add unit tests for AbstractTrainer save methods

---
 .../vsf_logging/MlflowLogger.py               | 10 +--
 .../test_abstract_trainer_save_artifacts.py   | 82 +++++++++++++++++++
 2 files changed, 84 insertions(+), 8 deletions(-)
 create mode 100644 tests/trainers/test_abstract_trainer_save_artifacts.py

diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
index c534378..62757f3 100644
--- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
+++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
@@ -276,16 +276,10 @@ def on_train_end(self):
         """
         # Save weights to a temporary directory and log artifacts
         if self._save_model_at_train_end:
-            self._save_model_weights(
-                artifact_path='weights',
-                best_model=False
-            )
+            _log_trainer_artifact(self.trainer, best_model=False)
 
         if self._save_best_model:
-            self._save_model_weights(
-                artifact_path='weights',
-                best_model=True
-            )
+            _log_trainer_artifact(self.trainer, best_model=True)
 
         for callback in self.callbacks:
             if hasattr(callback, 'on_train_end'):
diff --git a/tests/trainers/test_abstract_trainer_save_artifacts.py b/tests/trainers/test_abstract_trainer_save_artifacts.py
new file mode 100644
index 0000000..8f5ec49
--- /dev/null
+++ b/tests/trainers/test_abstract_trainer_save_artifacts.py
@@ -0,0 +1,82 @@
+"""Tests for AbstractTrainer save-related class methods."""
+
+import torch
+
+from virtual_stain_flow.trainers.AbstractTrainer import AbstractTrainer
+
+
+class MinimalArtifactTrainer(AbstractTrainer):
+    """Concrete trainer realization to exercise base save methods."""
+
+    def train_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict:
+        return {"loss": torch.tensor(0.0)}
+
+    def evaluate_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict:
+        return {"loss": torch.tensor(0.0)}
+
+
+def test_save_model_writes_current_model_file(
+    mock_model_with_save,
+    mock_optimizer,
+    train_dataloader,
+    val_dataloader,
+    tmp_path,
+):
+    trainer = MinimalArtifactTrainer(
+        model=mock_model_with_save,
+        optimizer=mock_optimizer,
+        train_loader=train_dataloader,
+        val_loader=val_dataloader,
+        device=torch.device("cpu"),
+    )
+
+    saved_paths = trainer.save_model(save_path=tmp_path, best_model=False)
+
+    assert isinstance(saved_paths, list)
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "generator_weights_0.pth"
+
+
+def test_save_optimizer_state_writes_recent_file(
+    mock_model_with_save,
+    mock_optimizer,
+    train_dataloader,
+    val_dataloader,
+    tmp_path,
+):
+    trainer = MinimalArtifactTrainer(
+        model=mock_model_with_save,
+        optimizer=mock_optimizer,
+        train_loader=train_dataloader,
+        val_loader=val_dataloader,
+        device=torch.device("cpu"),
+    )
+
+    saved_paths = trainer.save_optimizer_state(save_path=tmp_path, recent=True)
+
+    assert isinstance(saved_paths, list)
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "optimizer_recent.pth"
+
+
+def test_save_optimizer_state_non_recent_not_supported(
+    mock_model_with_save,
+    mock_optimizer,
+    train_dataloader,
+    val_dataloader,
+    tmp_path,
+):
+    trainer = MinimalArtifactTrainer(
+        model=mock_model_with_save,
+        optimizer=mock_optimizer,
+        train_loader=train_dataloader,
+        val_loader=val_dataloader,
+        device=torch.device("cpu"),
+    )
+
+    import pytest
+
+    with pytest.raises(NotImplementedError, match="non-recent optimizer states"):
+        trainer.save_optimizer_state(save_path=tmp_path, recent=False)

From 3b0d53b00a753387031416968ad9a1a05ebf228b Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 8 Sep 2026 23:29:04 -0600
Subject: [PATCH 05/10] Rename _log_trainer_artifact function name with
 _save_and_log_trainer_artifacts, added more detailed docstring, for improved
 naming clarity

---
 src/virtual_stain_flow/vsf_logging/MlflowLogger.py     | 10 +++++-----
 .../vsf_logging/logger_utils/__init__.py               |  4 ++--
 .../vsf_logging/logger_utils/log_artifacts.py          |  7 ++++---
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
index 62757f3..b58d023 100644
--- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
+++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
@@ -17,7 +17,7 @@
     AutoOptimizerConfigLogger,
     AutoTrainerLogger
 )
-from .logger_utils import _log_artifact, _log_trainer_artifact
+from .logger_utils import _log_artifact, _save_and_log_trainer_artifacts
 from .callbacks.LoggerCallback import (
     AbstractLoggerCallback,
     log_type
@@ -251,10 +251,10 @@ def on_epoch_end(self):
 
         if self._save_model_every_n_epochs is not None:
             if self.trainer.epoch % self._save_model_every_n_epochs == 0:
-                _log_trainer_artifact(self.trainer, best_model=False)
+                _save_and_log_trainer_artifacts(self.trainer, best_model=False)
 
         if self._save_best_model:
-            _log_trainer_artifact(self.trainer, best_model=True)
+            _save_and_log_trainer_artifacts(self.trainer, best_model=True)
 
         # Call on_epoch_end for all registered callbacks
         for callback in self.callbacks:
@@ -276,10 +276,10 @@ def on_train_end(self):
         """
         # Save weights to a temporary directory and log artifacts
         if self._save_model_at_train_end:
-            _log_trainer_artifact(self.trainer, best_model=False)
+            _save_and_log_trainer_artifacts(self.trainer, best_model=False)
 
         if self._save_best_model:
-            _log_trainer_artifact(self.trainer, best_model=True)
+            _save_and_log_trainer_artifacts(self.trainer, best_model=True)
 
         for callback in self.callbacks:
             if hasattr(callback, 'on_train_end'):
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
index 04ad9c8..e7edb5e 100644
--- a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
@@ -1,6 +1,6 @@
-from .log_artifacts import _log_artifact, _log_trainer_artifact
+from .log_artifacts import _log_artifact, _save_and_log_trainer_artifacts
 
 __all__ = [
     "_log_artifact",
-    "_log_trainer_artifact",
+    "_save_and_log_trainer_artifacts",
 ]
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
index 199fe7c..2bf371a 100644
--- a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
@@ -51,16 +51,17 @@ def _log_artifact(
     mlflow.log_artifact(str(file_path), artifact_path=artifact_path)
 
 
-def _log_trainer_artifact(
+def _save_and_log_trainer_artifacts(
     trainer: "TrainerProtocol",
     best_model: bool = True,
 ) -> None:
     """
-    Logs the model and optimizer state artifacts from a trainer to MLflow.
+    Saves the trainer's model and optimizer state to temporary files and logs
+    those files as MLflow artifacts.
     The most recent optimizer state is saved and logged regardless of the best_model flag.
 
     :param trainer: The trainer instance adhering to TrainerProtocol.
-    :param best_model: Whether to log only the best model, defaults to True.
+    :param best_model: Whether to save and log the best model, defaults to True.
     :raises TypeError: If the provided trainer does not adhere to TrainerProtocol.
     """
 

From f6248b93e25c2ac92d64c1e2ebf9ded049e0add4 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 9 Sep 2026 02:05:46 -0600
Subject: [PATCH 06/10] Update DatasetManifest serialization to stop using
 depracated functions, updating dependency as well.

---
 pyproject.toml                                |  2 +-
 .../datasets/ds_engine/input_validation.py    |  2 +-
 .../datasets/ds_engine/manifest.py            |  4 ++--
 tests/datasets/ds_engine/test_manifest.py     | 21 +++++++++++++++++++
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 56bd4fe..037c7fd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ requires-python = ">=3.9"
 dependencies = [
   "pyarrow",
   "numpy",
-  "pandas",
+  "pandas>=2.1",
   "seaborn",
   "matplotlib",
   "scikit-learn",
diff --git a/src/virtual_stain_flow/datasets/ds_engine/input_validation.py b/src/virtual_stain_flow/datasets/ds_engine/input_validation.py
index b55c38a..9bfaf62 100644
--- a/src/virtual_stain_flow/datasets/ds_engine/input_validation.py
+++ b/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.
diff --git a/src/virtual_stain_flow/datasets/ds_engine/manifest.py b/src/virtual_stain_flow/datasets/ds_engine/manifest.py
index 39b3884..36bb81a 100644
--- a/src/virtual_stain_flow/datasets/ds_engine/manifest.py
+++ b/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"""
diff --git a/tests/datasets/ds_engine/test_manifest.py b/tests/datasets/ds_engine/test_manifest.py
index 168071f..d174a78 100644
--- a/tests/datasets/ds_engine/test_manifest.py
+++ b/tests/datasets/ds_engine/test_manifest.py
@@ -362,6 +362,27 @@ def test_from_config_missing_manifest(self):
 class TestDatasetManifestSerialization:
     """Test suite for DatasetManifest serialization methods."""
 
+    def test_config_round_trip(self):
+        """Test serialization converts paths to strings and restores Paths."""
+        manifest = DatasetManifest(
+            file_index=pd.DataFrame({
+                "channel1": [Path("/path/to/img1.tif")],
+                "channel2": ["/path/to/img2.tif"],
+            })
+        )
+
+        config = manifest.to_config()
+        restored = DatasetManifest.from_config(config)
+
+        assert config["file_index"] == [{
+            "channel1": "/path/to/img1.tif",
+            "channel2": "/path/to/img2.tif",
+        }]
+        assert restored.file_index.to_dict(orient="records") == [{
+            "channel1": Path("/path/to/img1.tif"),
+            "channel2": Path("/path/to/img2.tif"),
+        }]
+
     def test_from_config_missing_file_index(self):
         """Test DatasetManifest.from_config raises ValueError when file_index is missing."""
         config = {"pil_image_mode": "I;16", "file_index": None}

From ac2c0b83a9d277310c25bcc76c985d4da64cb080 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 9 Sep 2026 02:12:08 -0600
Subject: [PATCH 07/10] Refactor protocol and abstract class definition to make
 the `file_name_prefix` parameter of the save_model and save_optimizer_state
 function non-optional, and directly assign default values in realization
 implementations. This reduces code and improves clarity.

---
 src/virtual_stain_flow/trainers/AbstractTrainer.py          | 6 +++---
 src/virtual_stain_flow/trainers/trainer_protocol.py         | 4 ++--
 src/virtual_stain_flow/trainers/trainer_utils/save_model.py | 5 +----
 .../trainers/trainer_utils/save_optimizer.py                | 5 +----
 tests/conftest.py                                           | 2 +-
 5 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 82bf209..00a89a6 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -412,7 +412,7 @@ 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
@@ -420,7 +420,7 @@ def save_model(
         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
@@ -429,7 +429,7 @@ def save_model(
     def save_optimizer_state(
         self, 
         save_path: pathlib.Path, 
-        file_name_prefix: Optional[str] = None, 
+        file_name_prefix: str = 'optimizer',
         file_name_suffix: Optional[str] = None, 
         file_ext: str = '.pth',
         recent: bool = True
diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index e0de924..da98198 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/src/virtual_stain_flow/trainers/trainer_protocol.py
@@ -71,7 +71,7 @@ def best_model(self) -> Optional[torch.nn.Module]: ...
     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,
@@ -81,7 +81,7 @@ def save_model(
     def save_optimizer_state(
         self, 
         save_path: pathlib.Path, 
-        file_name_prefix: Optional[str] = None, 
+        file_name_prefix: str = 'optimizer',
         file_name_suffix: Optional[str] = None, 
         file_ext: str = '.pth',
         recent: bool = True,
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
index 8063199..528d81b 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
@@ -6,15 +6,12 @@
 def save_model(
     trainer: 'TrainerProtocol',
     save_path: Path,
-    file_name_prefix: Optional[str] = None,
+    file_name_prefix: str = 'generator',
     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)
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
index 97e899c..ecdb510 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
@@ -9,14 +9,11 @@
 def save_optimizer_state(
     trainer: 'TrainerProtocol',
     save_path: Path,
-    file_name_prefix: Optional[str] = None,
+    file_name_prefix: str = 'optimizer',
     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}"
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 5842619..9c42267 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -388,7 +388,7 @@ def evaluate_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict:
             'loss_b': torch.tensor(0.2),
         }
 
-    def save_model(self, save_path, file_name_prefix=None, file_name_suffix=None,
+    def save_model(self, save_path, file_name_prefix='generator', file_name_suffix=None,
                    file_ext='.pth', best_model=True):
         return None
 

From cc3ab85d6ce5275debd6103c70e4b0ff5d0a4e11 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 9 Sep 2026 13:43:38 -0600
Subject: [PATCH 08/10] Remove unused _epochs attribute from TrainerProtocol
 class

---
 src/virtual_stain_flow/trainers/trainer_protocol.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index da98198..65a5143 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/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
 

From 309927acb3f7e9a77d079c170d6caec72c448900 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 9 Sep 2026 14:12:43 -0600
Subject: [PATCH 09/10] Refactor AbstractTrainer to enforce integer type for
 epoch parameter and add validation. Introduce utility method for dataset size
 retrieval and enhance tests for epoch handling and dataset size scenarios.

---
 .../trainers/AbstractTrainer.py               | 37 ++++++++++++++-----
 tests/trainers/test_abstract_trainer.py       | 31 ++++++++++++++++
 tests/trainers/test_logging_trainer.py        | 14 +++++--
 3 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 00a89a6..f45741a 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -42,7 +42,7 @@ def __init__(
         test_ratio: Optional[float] = 0.15,
         metrics: Dict[str, AbstractMetrics] = None,
         device: Optional[torch.device] = None,
-        epoch: Optional[int] = 0,
+        epoch: int = 0,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
@@ -69,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
@@ -108,17 +110,20 @@ def __init__(
 
     def _init_state(
         self, 
-        epoch,
+        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 = 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)
@@ -197,12 +202,23 @@ def _init_data(
             )
 
         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
+        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]:
         """
@@ -330,16 +346,17 @@ def train(
         if hasattr(logger, "on_train_start"):
             logger.on_train_start()
 
+        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"):
diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py
index 0d22901..4474603 100644
--- a/tests/trainers/test_abstract_trainer.py
+++ b/tests/trainers/test_abstract_trainer.py
@@ -338,6 +338,37 @@ def test_train_epoch_with_large_batch_count(self, minimal_model, minimal_optimiz
 
 class TestDataSplitting:
     """Test that AbstractTrainer correctly handles dataset splitting."""
+
+    def test_init_rejects_none_epoch(
+        self, minimal_model, minimal_optimizer, train_dataloader
+    ):
+        with pytest.raises(TypeError, match="epoch must be an integer"):
+            MinimalTrainerRealization(
+                model=minimal_model,
+                optimizer=minimal_optimizer,
+                train_loader=train_dataloader,
+                epoch=None,
+                device=torch.device('cpu')
+            )
+
+    def test_init_with_unsized_dataset_records_unknown_size(
+        self, minimal_model, minimal_optimizer
+    ):
+        from torch.utils.data import DataLoader, IterableDataset
+
+        class UnsizedDataset(IterableDataset):
+            def __iter__(self):
+                yield torch.randn(4), torch.randn(2)
+
+        train_loader = DataLoader(UnsizedDataset(), batch_size=1)
+        trainer = MinimalTrainerRealization(
+            model=minimal_model,
+            optimizer=minimal_optimizer,
+            train_loader=train_loader,
+            device=torch.device('cpu')
+        )
+
+        assert trainer.train_n is None
     
     def test_init_with_dataset_creates_loaders(self, minimal_model, minimal_optimizer, dataset_for_splitting):
         """Verify that providing a dataset creates train/val/test loaders."""
diff --git a/tests/trainers/test_logging_trainer.py b/tests/trainers/test_logging_trainer.py
index c4c455b..0250272 100644
--- a/tests/trainers/test_logging_trainer.py
+++ b/tests/trainers/test_logging_trainer.py
@@ -559,6 +559,14 @@ def test_train_logged_metrics_have_correct_steps(self, conv_trainer, dummy_logge
         epochs = 3
         conv_trainer.train(logger=dummy_logger, epochs=epochs, verbose=False)
         
-        # Check that step numbers are within expected range (0-indexed)
-        for metric in dummy_logger.logged_metrics:
-            assert 0 <= metric['step'] < epochs
+        assert {metric['step'] for metric in dummy_logger.logged_metrics} == {1, 2, 3}
+
+    def test_resumed_training_logs_cumulative_epoch_steps(
+        self, conv_trainer, dummy_logger
+    ):
+        """Metric steps stay aligned with callbacks after resumed training."""
+        conv_trainer.epoch = 5
+
+        conv_trainer.train(logger=dummy_logger, epochs=2, verbose=False)
+
+        assert {metric['step'] for metric in dummy_logger.logged_metrics} == {6, 7}

From 6b13039a14c4733e8506f0e349911a48d2ca79e2 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Fri, 11 Sep 2026 14:49:59 -0600
Subject: [PATCH 10/10] Update CHANGELOG.md for version 0.4.9: Add initial
 training resume support and fix deprecated Pandas API usage

---
 CHANGELOG.md   | 15 +++++++++++++++
 pyproject.toml |  2 +-
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4790da6..1de026b 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index 037c7fd..a9497a2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ 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 = [