From d0ff8ce12b41c3002f60ef637f170da2fd3636f8 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Mon, 3 Aug 2026 17:09:05 -0600
Subject: [PATCH 1/8] Add stand-alone trainer utility modules for early
 stopping and model saving, in preparation for refactoring of
 `AbstractTrainer`.

---
 .../trainers/trainer_utils/__init__.py        |  18 +++
 .../trainers/trainer_utils/early_stop.py      | 147 ++++++++++++++++++
 .../trainers/trainer_utils/save_model.py      |  30 ++++
 3 files changed, 195 insertions(+)
 create mode 100644 src/virtual_stain_flow/trainers/trainer_utils/__init__.py
 create mode 100644 src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
 create mode 100644 src/virtual_stain_flow/trainers/trainer_utils/save_model.py

diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
new file mode 100644
index 0000000..d106587
--- /dev/null
+++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
@@ -0,0 +1,18 @@
+"""
+Trainer utilities for virtual stain flow.
+"""
+
+from .early_stop import (
+    BestEpochState,
+    EarlyStopHelper,
+    _get_latest_metric_value,
+)
+from .save_model import save_model
+
+
+__all__ = [
+    "BestEpochState",
+    "EarlyStopHelper",
+    "_get_latest_metric_value",
+    "save_model",
+]
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py b/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
new file mode 100644
index 0000000..a2a6841
--- /dev/null
+++ b/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
@@ -0,0 +1,147 @@
+from dataclasses import dataclass, field
+from typing import Literal, Optional
+
+import torch
+
+from ...models.model import clone_model
+
+
+@dataclass
+class BestEpochState:
+    """Track the best validation observation and its model snapshot."""
+
+    best_mode: Literal["min", "max"] = "min"
+    best_metric_value: Optional[float] = None
+    best_epoch: int = 0
+    best_model: Optional[torch.nn.Module] = field(default=None, init=False)
+
+    def __post_init__(self) -> None:
+        if self.best_mode not in ("min", "max"):
+            raise ValueError("best_mode must be either 'min' or 'max'.")
+
+        if self.best_metric_value is None:
+            self.best_metric_value = (
+                float("inf") if self.best_mode == "min" else float("-inf")
+            )
+
+    def update(
+        self,
+        metric_value: float,
+        epoch: Optional[int] = None,
+        model: Optional[torch.nn.Module] = None,
+    ) -> bool:
+        """Record an improved observation and return whether it improved."""
+        metric_value = float(metric_value)
+        is_better = (
+            self.best_mode == "min" and metric_value < self.best_metric_value
+        ) or (
+            self.best_mode == "max" and metric_value > self.best_metric_value
+        )
+        if not is_better:
+            return False
+
+        self.best_metric_value = metric_value
+        self.best_epoch = epoch if epoch is not None else self.best_epoch + 1
+        if model is not None:
+            self.best_model = clone_model(model)
+
+        return True
+
+
+class EarlyStopHelper:
+    """Track validation improvements and decide when training should stop."""
+
+    def __init__(
+        self,
+        best_mode: Literal["min", "max"] = "min",
+        trainer_val_losses_ref: Optional[dict] = None,
+        trainer_val_metrics_ref: Optional[dict] = None,
+        best_metric_name: Optional[str] = None,
+        enabled: bool = True,
+    ) -> None:
+        self.best_epoch_state = BestEpochState(best_mode=best_mode)
+        self.trainer_val_losses_ref = (
+            trainer_val_losses_ref if trainer_val_losses_ref is not None else {}
+        )
+        self.trainer_val_metrics_ref = (
+            trainer_val_metrics_ref if trainer_val_metrics_ref is not None else {}
+        )
+        self.best_metric_name = best_metric_name
+        self.enabled = enabled
+        self._patience: Optional[int] = None
+        self._counter = 0
+
+    def initialize_early_stop(self, patience: int) -> None:
+        """Initialize a run's patience counter while preserving the best state."""
+        if patience < 1:
+            raise ValueError("patience must be at least 1.")
+
+        self._patience = patience
+        self._counter = 0
+
+    def update(
+        self,
+        epoch: Optional[int] = None,
+        model: Optional[torch.nn.Module] = None,
+    ) -> bool:
+        """Update tracking state and return whether training should stop."""
+        if not self.enabled:
+            return False
+        if self._patience is None:
+            raise RuntimeError(
+                "Early stopping has not been initialized. "
+                "Call 'initialize_early_stop' first."
+            )
+
+        metric_value = _get_latest_metric_value(
+            self.trainer_val_losses_ref,
+            self.trainer_val_metrics_ref,
+            self.best_metric_name,
+        )
+        if self.best_epoch_state.update(metric_value, epoch, model):
+            self._counter = 0
+        else:
+            self._counter += 1
+
+        return self._counter >= self._patience
+
+    @property
+    def best_model(self) -> Optional[torch.nn.Module]:
+        """Return the best model snapshot, if an observation was recorded."""
+        return self.best_epoch_state.best_model
+
+    @property
+    def counter(self) -> int:
+        """Return the number of consecutive observations without improvement."""
+        return self._counter
+
+
+def _get_latest_metric_value(
+    val_losses: dict,
+    val_metrics: dict,
+    metric_name: Optional[str] = None,
+) -> float:
+    """Return the latest configured validation loss or metric value."""
+    if metric_name is None:
+        if not val_losses:
+            raise RuntimeError("No validation losses are available for early stopping.")
+        metric_name = next(iter(val_losses))
+
+    if metric_name in val_losses:
+        values = val_losses[metric_name]
+    elif metric_name in val_metrics:
+        values = val_metrics[metric_name]
+    else:
+        raise ValueError(
+            f"Supplied early stop metric '{metric_name}' is not found in logs."
+        )
+
+    if not isinstance(values, (list, tuple)):
+        raise TypeError(
+            f"Expected the metric '{metric_name}' to be a list or tuple, "
+            f"but got {type(values)} instead."
+        )
+    if not values:
+        raise RuntimeError(f"Early stop metric '{metric_name}' has no recorded values.")
+
+    return float(values[-1])
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
new file mode 100644
index 0000000..528d81b
--- /dev/null
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
@@ -0,0 +1,30 @@
+from pathlib import Path
+from typing import Optional, List
+
+from ..trainer_protocol import TrainerProtocol
+
+def save_model(
+    trainer: 'TrainerProtocol',
+    save_path: Path,
+    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_suffix is None:
+        file_name_suffix = 'weights_' + (
+            'best' if save_best_model else str(trainer.epoch)
+        )
+
+    model = trainer.best_model if save_best_model else trainer.model
+
+    if model is None:
+        return []
+
+    path = model.save_weights(
+        filename=f"{file_name_prefix}_{file_name_suffix}{file_ext}",
+        dir=save_path
+    )
+
+    return [path]

From 5cc5d78b52b7e310a1b672e7415430b9b22c6e76 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Mon, 3 Aug 2026 17:09:16 -0600
Subject: [PATCH 2/8] Add deep copy functionality to BaseModel and implement
 clone_model utility

---
 src/virtual_stain_flow/models/model.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/virtual_stain_flow/models/model.py b/src/virtual_stain_flow/models/model.py
index b2acaf8..7dd87f3 100644
--- a/src/virtual_stain_flow/models/model.py
+++ b/src/virtual_stain_flow/models/model.py
@@ -16,6 +16,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Dict, Union, Any
 import pathlib
+import copy
 
 import torch
 
@@ -97,6 +98,15 @@ def from_config(cls, config: Dict) -> 'BaseModel':
         """
         raise NotImplementedError("Subclasses must implement this method.")
 
+    def copy(self) -> 'BaseModel':
+        """
+        Creates a deep copy of the model instance.
+        
+        :return: A deep copy of the model.
+        """
+        return clone_model(self)
+    
+
 class BaseGeneratorModel(BaseModel):
 
     def __init__(
@@ -148,3 +158,18 @@ def out_channels(self) -> int:
     def out_activation(self) -> torch.nn.Module:
         """Activation function for the output layer."""
         return self._out_activation
+
+
+def clone_model(model: torch.nn.Module) -> torch.nn.Module:
+    try:
+        snapshot = copy.deepcopy(model)
+    except Exception as error:
+        raise RuntimeError(
+            f"Unable to create an independent snapshot of {type(model).__name__}."
+        ) from error
+
+    if not isinstance(snapshot, torch.nn.Module) or snapshot is model:
+        raise RuntimeError(
+            f"Unable to create an independent snapshot of {type(model).__name__}."
+        )
+    return snapshot

From cd73fe2ec5d9328fd607fe12fab7a826b04b6129 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Mon, 3 Aug 2026 17:14:42 -0600
Subject: [PATCH 3/8] Refactor `AbstractTrainer` to wire early stopping related
 logic to EarlyStopHelper remove internal logic. Remove `AbstractTrainer`
 attributes that are not used by any other components of this software and no
 longer updated by the most recent trainer implementation. Centralize save
 model implementation as previous realizations across the two trainer classes
 are effectively identical. Add tests to contract new trainer model saving
 behavior when validiation set is not present.

---
 .../trainers/AbstractTrainer.py               | 119 ++++--------
 .../trainers/logging_gan_trainer.py           |  21 --
 .../trainers/logging_trainer.py               |  25 ---
 .../trainers/trainer_protocol.py              |   2 +-
 tests/trainers/test_abstract_trainer.py       | 181 ++----------------
 tests/trainers/test_logging_trainer.py        |  61 +++++-
 .../trainers/trainer_utils/test_early_stop.py | 164 ++++++++++++++++
 7 files changed, 267 insertions(+), 306 deletions(-)
 create mode 100644 tests/trainers/trainer_utils/test_early_stop.py

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 36c25ae..a1077e5 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -13,6 +13,7 @@
 from torch.utils.data import DataLoader
 
 from .trainer_protocol import TrainerProtocol
+from .trainer_utils import EarlyStopHelper, save_model
 from ..metrics.AbstractMetrics import AbstractMetrics
 from ..engine.progress import Progress
 from ..datasets.data_split import default_random_split
@@ -41,7 +42,7 @@ def __init__(
         test_ratio: Optional[float] = 0.15,
         metrics: Dict[str, AbstractMetrics] = None,
         device: Optional[torch.device] = None,
-        early_termination_metric: str = None,
+        early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
     ):
@@ -100,17 +101,10 @@ def __init__(
 
     def _init_state(
         self, 
-        early_termination_metric,
-        early_termination_mode,
+        early_termination_metric: Optional[str] = None,
+        early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs
     ):
-        # Early stopping state
-        self._best_model = None
-        self._best_loss = float("inf")
-        self._early_stop_counter = 0
-        self._early_termination_metric = early_termination_metric
-        self._early_termination_mode = early_termination_mode
-        self._early_termination = True if early_termination_metric else False
 
         # Epoch state
         self._epoch = 0
@@ -124,6 +118,15 @@ def _init_state(
         self._train_metrics = defaultdict(list)
         self._val_metrics = defaultdict(list)
 
+        # Early stopping state
+        self._early_stop_helper = EarlyStopHelper(
+            best_mode=early_termination_mode,
+            trainer_val_losses_ref=self._val_losses,
+            trainer_val_metrics_ref=self._val_metrics,
+            best_metric_name=early_termination_metric,
+            enabled=bool(self._val_loader),
+        )
+
         return None
 
     def _init_data(
@@ -308,11 +311,12 @@ def train(
             logger.on_train_start()
 
         self._epochs = epochs
-        self._patience = patience if patience else epochs # no early stopping
         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)
 
+        self._early_stop_helper.initialize_early_stop(patience=patience if patience else epochs)
+
         for epoch in iterable:
 
             # Increment the epoch counter
@@ -355,65 +359,22 @@ def train(
                     f"val_{metric_name}", val_metric, epoch
                 )
 
+            # Update early stopping
+            should_stop = self._early_stop_helper.update(
+                epoch=self.epoch,
+                model=self.model
+            )
+
             if hasattr(logger, "on_epoch_end"):
                 logger.on_epoch_end()
 
-            # Update early stopping
-            if self.update_early_stop_counter():
-                print(f"Early termination at epoch {epoch + 1} "
-                      f"with best validation metric {self._best_loss}")
+            if should_stop:
+                print(f"Early termination at epoch {self.epoch} "
+                      f"with best validation metric {self._early_stop_helper.best_epoch_state.best_metric_value}")
                 break
 
         if hasattr(logger, "on_train_end"):
             logger.on_train_end()
-
-    def _collect_early_stop_metric(self) -> Optional[float]:
-        if self._early_termination_metric is None:
-            # Do not perform early stopping when no termination metric is specified
-            early_term_metric = None
-        else:
-            # First look for the metric in validation loss
-            if self._early_termination_metric in list(
-                self._val_losses.keys()):
-                early_term_metric = self._val_losses[
-                    self._early_termination_metric][-1]
-            # Then look for the metric in validation metrics
-            elif self._early_termination_metric in list(
-                self._val_metrics.keys()):
-                early_term_metric = self._val_metrics[
-                    self._early_termination_metric][-1]
-            else:
-                raise ValueError("Invalid early termination metric")
-            
-        return early_term_metric
-
-    def update_early_stop_counter(self) -> bool:
-        """
-        Method to update the early stopping criterion
-
-        :return: True if early stopping criterion is met, False otherwise.
-        """
-        
-        early_term_metric = self._collect_early_stop_metric()
-
-        # When early termination is disabled, 
-        # the best model is updated with the current model
-        if not self._early_termination and early_term_metric is None:
-            self.best_model = self.model.state_dict().copy()
-            return False
-        
-        reset_counter = (early_term_metric < self.best_loss) \
-            if self._early_termination_mode == "min" \
-                else (early_term_metric > self.best_loss)
-
-        if reset_counter:
-            self.best_loss = early_term_metric
-            self.early_stop_counter = 0
-            self.best_model = self.model.state_dict().copy()
-        else:
-            self.early_stop_counter += 1
-
-        return self.early_stop_counter >= self.patience
     
     def _update_epoch_progress(
         self,
@@ -430,7 +391,6 @@ def _update_epoch_progress(
             f"{phase} Batch {batch_idx + 1}/{num_batches}"
         )
 
-    @abstractmethod
     def save_model(
         self,
         save_path: pathlib.Path,
@@ -439,7 +399,14 @@ def save_model(
         file_ext: str = '.pth',
         best_model: bool = True
     ) -> Optional[List[pathlib.Path]]:
-        pass
+        return save_model(
+            self,
+            save_path=save_path,
+            file_name_prefix=file_name_prefix or 'generator',
+            file_name_suffix=file_name_suffix,
+            file_ext=file_ext,
+            save_best_model=best_model
+        )
 
     """
     Log property
@@ -500,15 +467,7 @@ def patience(self):
         
     @property
     def best_model(self):
-        return self._best_model
-    
-    @property
-    def best_loss(self):
-        return self._best_loss
-    
-    @property
-    def early_stop_counter(self):
-        return self._early_stop_counter
+        return self._early_stop_helper.best_model
     
     @property
     def metrics(self):
@@ -544,18 +503,6 @@ def val_metrics(self):
     Meant to be used by the subclasses to update the best model and loss
     """
 
-    @best_model.setter
-    def best_model(self, value: torch.nn.Module):
-        self._best_model = value
-    
-    @best_loss.setter
-    def best_loss(self, value):
-        self._best_loss = value
-
-    @early_stop_counter.setter
-    def early_stop_counter(self, value: int):
-        self._early_stop_counter = value
-
     @epoch.setter
     def epoch(self, value: int):
         self._epoch = value
diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
index d959457..11a6d8c 100644
--- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
@@ -6,7 +6,6 @@
 model using the engine subpackage for forward passes and loss computations.
 """
 
-import pathlib
 from typing import Dict, List, Union, Optional
 
 import torch
@@ -200,27 +199,7 @@ def loss_groups(self) -> Dict[str, LossGroup]:
             'generator': self._generator_loss_group,
             'discriminator': self._discriminator_loss_group
         }
-    
-    def save_model(
-        self,
-        save_path: pathlib.Path,
-        file_name_prefix: Optional[str] = None,
-        file_name_suffix: Optional[str] = None,
-        file_ext: str = '.pth',
-        best_model: bool = True
-    ) -> Optional[List[pathlib.Path]]:
-
-        if file_name_suffix is None:
-            file_name_suffix = 'weights_' + (
-                'best' if best_model else str(self.epoch)
-            )
-
-        gen_path = self.model.save_weights(
-            filename=f"generator_{file_name_suffix}{file_ext}",
-            dir=save_path
-        )
 
-        return [gen_path]
 
 class LoggingWGANTrainer(BaseGANTrainer):
     """
diff --git a/src/virtual_stain_flow/trainers/logging_trainer.py b/src/virtual_stain_flow/trainers/logging_trainer.py
index ff20c8d..645910f 100644
--- a/src/virtual_stain_flow/trainers/logging_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_trainer.py
@@ -7,7 +7,6 @@
 computations.
 """
 
-import pathlib
 from typing import Dict, List, Union, Optional
 
 import torch
@@ -151,27 +150,3 @@ def evaluate_step(
     @property
     def loss_groups(self) -> Dict[str, LossGroup]:
         return {'main': self._loss_group}
-    
-    def save_model(
-        self,
-        save_path: pathlib.Path,
-        file_name_prefix: Optional[str] = None,
-        file_name_suffix: Optional[str] = None,
-        file_ext: str = '.pth',
-        best_model: bool = True
-    ) -> Optional[List[pathlib.Path]]:
-
-        if file_name_prefix is None:
-            file_name_prefix = 'generator'
-
-        if file_name_suffix is None:
-            file_name_suffix = 'weights_' + (
-                'best' if best_model else str(self.epoch)
-            )
-
-        path = self.model.save_weights(
-            filename=f"{file_name_prefix}_{file_name_suffix}{file_ext}",
-            dir=save_path
-        )
-
-        return [path]
diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index 80b55cc..5fc3cb0 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/src/virtual_stain_flow/trainers/trainer_protocol.py
@@ -54,7 +54,7 @@ def metrics(self) -> Dict[str, torch.nn.Module]: ...
     def model(self) -> torch.nn.Module: ...
 
     @property
-    def best_model(self) -> torch.nn.Module: ...
+    def best_model(self) -> Optional[torch.nn.Module]: ...
 
     def save_model(
         self, 
diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py
index 62d6e38..2cb9099 100644
--- a/tests/trainers/test_abstract_trainer.py
+++ b/tests/trainers/test_abstract_trainer.py
@@ -202,6 +202,17 @@ def test_epoch_counter_not_incremented_by_evaluate_epoch(self, trainer_with_load
 
 class TestTrainEpochEdgeCases:
     """Test edge cases for train_epoch and evaluate_epoch."""
+
+    def test_train_without_validation_completes_all_epochs(
+        self, trainer_with_empty_val_loader, dummy_logger
+    ):
+        trainer = trainer_with_empty_val_loader
+
+        trainer.train(logger=dummy_logger, epochs=3, patience=1, verbose=False)
+
+        assert trainer.epoch == 3
+        assert trainer.best_model is None
+        assert all(len(values) == 3 for values in trainer.train_losses.values())
     
     def test_train_epoch_with_empty_dataloader(self, minimal_model, minimal_optimizer, empty_dataloader):
         """
@@ -528,177 +539,7 @@ def test_log_property_with_no_data(self, trainer_with_loaders):
         assert "epoch" in log
         assert log["epoch"] == []
 
-
-class TestEarlyTermination:
-    """Test early termination logic."""
-    
-    def test_early_termination_disabled_by_default(self, trainer_with_loaders):
-        """Verify that early termination is disabled when no metric is specified."""
-        trainer = trainer_with_loaders
-        
-        assert trainer._early_termination is False
-        assert trainer._early_termination_metric is None
-    
-    def test_early_termination_enabled_with_metric(self, minimal_model, minimal_optimizer, train_dataloader, val_dataloader):
-        """Verify that early termination is enabled when metric is specified."""
-        trainer = MinimalTrainerRealization(
-            model=minimal_model,
-            optimizer=minimal_optimizer,
-            train_loader=train_dataloader,
-            val_loader=val_dataloader,
-            batch_size=2,
-            device=torch.device('cpu'),
-            early_termination_metric='loss_a'
-        )
-        
-        assert trainer._early_termination is True
-        assert trainer._early_termination_metric == 'loss_a'
-    
-    def test_update_early_stop_counter_improves(self, minimal_model, minimal_optimizer, train_dataloader, val_dataloader):
-        """Verify that early stop counter resets when validation loss improves."""
-        trainer = MinimalTrainerRealization(
-            model=minimal_model,
-            optimizer=minimal_optimizer,
-            train_loader=train_dataloader,
-            val_loader=val_dataloader,
-            batch_size=2,
-            device=torch.device('cpu'),
-            early_termination_metric='val_loss',
-            early_termination_mode='min'
-        )
-        
-        trainer._patience = 3
-        trainer._early_stop_counter = 2
-        trainer._best_loss = 0.5
-        
-        # Add a better validation loss
-        trainer.update_loss(torch.tensor(0.3), "val_loss", validation=True)
-        
-        should_stop = trainer.update_early_stop_counter()
-        
-        assert trainer.early_stop_counter == 0  # Reset
-        assert trainer.best_loss == 0.3  # Updated
-        assert should_stop is False
-    
-    def test_update_early_stop_counter_worsens(self, minimal_model, minimal_optimizer, train_dataloader, val_dataloader):
-        """Verify that early stop counter increments when validation loss worsens."""
-        trainer = MinimalTrainerRealization(
-            model=minimal_model,
-            optimizer=minimal_optimizer,
-            train_loader=train_dataloader,
-            val_loader=val_dataloader,
-            batch_size=2,
-            device=torch.device('cpu'),
-            early_termination_metric='val_loss',
-            early_termination_mode='min'
-        )
-        
-        trainer._patience = 3
-        trainer._early_stop_counter = 1
-        trainer._best_loss = 0.3
-        
-        # Add a worse validation loss
-        trainer.update_loss(torch.tensor(0.5), "val_loss", validation=True)
-        
-        should_stop = trainer.update_early_stop_counter()
-        
-        assert trainer.early_stop_counter == 2  # Incremented
-        assert trainer.best_loss == 0.3  # Not updated
-        assert should_stop is False
-    
-    def test_update_early_stop_counter_triggers_stop(self, minimal_model, minimal_optimizer, train_dataloader, val_dataloader):
-        """Verify that early stopping triggers when patience is exceeded."""
-        trainer = MinimalTrainerRealization(
-            model=minimal_model,
-            optimizer=minimal_optimizer,
-            train_loader=train_dataloader,
-            val_loader=val_dataloader,
-            batch_size=2,
-            device=torch.device('cpu'),
-            early_termination_metric='val_loss',
-            early_termination_mode='min'
-        )
-        
-        trainer._patience = 3
-        trainer._early_stop_counter = 2
-        trainer._best_loss = 0.3
-        
-        # Add a worse validation loss
-        trainer.update_loss(torch.tensor(0.5), "val_loss", validation=True)
-        
-        should_stop = trainer.update_early_stop_counter()
-        
-        assert trainer.early_stop_counter == 3
-        assert should_stop is True
     
-    def test_early_termination_mode_max(self, minimal_model, minimal_optimizer, train_dataloader, val_dataloader):
-        """Verify that early termination works in 'max' mode (e.g., for accuracy)."""
-        trainer = MinimalTrainerRealization(
-            model=minimal_model,
-            optimizer=minimal_optimizer,
-            train_loader=train_dataloader,
-            val_loader=val_dataloader,
-            batch_size=2,
-            device=torch.device('cpu'),
-            early_termination_metric='accuracy',
-            early_termination_mode='max'
-        )
-        
-        trainer._patience = 3
-        trainer._early_stop_counter = 1
-        trainer._best_loss = 0.8
-        
-        # Add a better accuracy (higher is better)
-        trainer.update_metrics(torch.tensor(0.9), "accuracy", validation=True)
-        
-        should_stop = trainer.update_early_stop_counter()
-        
-        assert trainer.early_stop_counter == 0  # Reset
-        assert trainer.best_loss == 0.9  # Updated
-        assert should_stop is False
-    
-    def test_collect_early_stop_metric_from_val_losses(self, trainer_with_loaders):
-        """Verify that early stop metric is collected from val_losses."""
-        trainer = trainer_with_loaders
-        trainer._early_termination_metric = "mse_loss"
-        
-        trainer.update_loss(torch.tensor(0.5), "mse_loss", validation=True)
-        
-        metric = trainer._collect_early_stop_metric()
-        
-        assert metric == torch.tensor(0.5)
-    
-    def test_collect_early_stop_metric_from_val_metrics(self, trainer_with_loaders):
-        """Verify that early stop metric is collected from val_metrics."""
-        trainer = trainer_with_loaders
-        trainer._early_termination_metric = "accuracy"
-        
-        trainer.update_metrics(torch.tensor(0.85), "accuracy", validation=True)
-        
-        metric = trainer._collect_early_stop_metric()
-        
-        assert metric == torch.tensor(0.85)
-    
-    def test_collect_early_stop_metric_invalid_metric_raises_error(self, trainer_with_loaders):
-        """Verify that invalid early termination metric raises ValueError."""
-        trainer = trainer_with_loaders
-        trainer._early_termination_metric = "nonexistent_metric"
-        
-        with pytest.raises(ValueError, match="Invalid early termination metric"):
-            trainer._collect_early_stop_metric()
-    
-    def test_early_termination_disabled_updates_best_model(self, trainer_with_loaders):
-        """Verify that when early termination is disabled, best model is still updated."""
-        trainer = trainer_with_loaders
-        trainer._early_termination = False
-        trainer._early_termination_metric = None
-        
-        should_stop = trainer.update_early_stop_counter()
-        
-        assert trainer.best_model is not None
-        assert should_stop is False
-
-
 class TestProperties:
     """Test that AbstractTrainer dataset properties work correctly."""
 
diff --git a/tests/trainers/test_logging_trainer.py b/tests/trainers/test_logging_trainer.py
index e6228d9..ff1f44d 100644
--- a/tests/trainers/test_logging_trainer.py
+++ b/tests/trainers/test_logging_trainer.py
@@ -186,11 +186,11 @@ def test_init_loss_items_on_correct_device(
 class TestSingleGeneratorTrainerSaveModel:
     """Tests for SingleGeneratorTrainer.save_model method."""
 
-    def test_save_model_creates_file(
+    def test_save_current_model_creates_file(
         self, mock_model_with_save, mock_optimizer, simple_loss, 
         train_dataloader, val_dataloader
     ):
-        """Test that save_model creates a file."""
+        """Test that saving the current model creates a file."""
         from virtual_stain_flow.trainers.logging_trainer import SingleGeneratorTrainer
         
         trainer = SingleGeneratorTrainer(
@@ -206,12 +206,67 @@ def test_save_model_creates_file(
         with tempfile.TemporaryDirectory() as tmpdir:
             tmpdir_path = pathlib.Path(tmpdir)
             
-            paths = trainer.save_model(save_path=tmpdir_path, best_model=True)
+            paths = trainer.save_model(save_path=tmpdir_path, best_model=False)
             
             assert paths is not None
             assert len(paths) == 1
             assert paths[0].exists()
 
+    def test_save_best_model_without_snapshot_returns_empty_list(
+        self, mock_model_with_save, mock_optimizer, simple_loss,
+        train_dataloader, val_dataloader
+    ):
+        """Test that saving a best model requires an established best snapshot."""
+        from virtual_stain_flow.trainers.logging_trainer import SingleGeneratorTrainer
+
+        trainer = SingleGeneratorTrainer(
+            model=mock_model_with_save,
+            optimizer=mock_optimizer,
+            losses=simple_loss,
+            device=torch.device('cpu'),
+            train_loader=train_dataloader,
+            val_loader=val_dataloader,
+            batch_size=2
+        )
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = trainer.save_model(
+                save_path=pathlib.Path(tmpdir),
+                best_model=True
+            )
+
+            assert paths == []
+
+    def test_save_best_model_creates_file_after_best_is_established(
+        self, mock_model_with_save, mock_optimizer, simple_loss,
+        train_dataloader, val_dataloader
+    ):
+        """Test that an established best-model snapshot is saved."""
+        from virtual_stain_flow.trainers.logging_trainer import SingleGeneratorTrainer
+
+        trainer = SingleGeneratorTrainer(
+            model=mock_model_with_save,
+            optimizer=mock_optimizer,
+            losses=simple_loss,
+            device=torch.device('cpu'),
+            train_loader=train_dataloader,
+            val_loader=val_dataloader,
+            batch_size=2
+        )
+        trainer._early_stop_helper.initialize_early_stop(patience=1)
+        trainer.val_losses["MSELoss"].append(0.5)
+        trainer._early_stop_helper.update(epoch=1, model=trainer.model)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = trainer.save_model(
+                save_path=pathlib.Path(tmpdir),
+                best_model=True
+            )
+
+            assert len(paths) == 1
+            assert paths[0].exists()
+            assert paths[0].name == "generator_weights_best.pth"
+
     def test_save_model_returns_list_of_paths(
         self, mock_model_with_save, mock_optimizer, simple_loss, 
         train_dataloader, val_dataloader
diff --git a/tests/trainers/trainer_utils/test_early_stop.py b/tests/trainers/trainer_utils/test_early_stop.py
new file mode 100644
index 0000000..941f848
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_early_stop.py
@@ -0,0 +1,164 @@
+import pytest
+import torch
+
+from virtual_stain_flow.trainers.trainer_utils.early_stop import (
+    BestEpochState,
+    EarlyStopHelper,
+    _get_latest_metric_value,
+)
+
+
+class TestBestEpochState:
+    @pytest.mark.parametrize(
+        ("best_mode", "initial_value", "new_value"),
+        [("min", 0.5, 0.3), ("max", 0.8, 0.9)],
+    )
+    def test_update_records_improvement(
+        self, best_mode, initial_value, new_value
+    ):
+        state = BestEpochState(
+            best_mode=best_mode,
+            best_metric_value=initial_value,
+        )
+
+        is_better = state.update(new_value, epoch=4)
+
+        assert is_better is True
+        assert state.best_metric_value == new_value
+        assert state.best_epoch == 4
+
+    def test_update_ignores_worse_value(self):
+        state = BestEpochState(best_metric_value=0.3, best_epoch=2)
+
+        is_better = state.update(0.5, epoch=3)
+
+        assert is_better is False
+        assert state.best_metric_value == 0.3
+        assert state.best_epoch == 2
+
+    def test_update_snapshots_model(self, minimal_model):
+        state = BestEpochState()
+
+        state.update(0.5, model=minimal_model)
+
+        assert state.best_model is not None
+        assert state.best_model is not minimal_model
+        for copied_parameter, parameter in zip(
+            state.best_model.parameters(), minimal_model.parameters()
+        ):
+            assert torch.equal(copied_parameter, parameter)
+
+
+class TestGetLatestMetricValue:
+    def test_gets_value_from_losses(self):
+        value = _get_latest_metric_value(
+            val_losses={"mse_loss": [0.5]},
+            val_metrics={},
+            metric_name="mse_loss",
+        )
+
+        assert value == 0.5
+
+    def test_gets_value_from_metrics(self):
+        value = _get_latest_metric_value(
+            val_losses={},
+            val_metrics={"accuracy": [0.85]},
+            metric_name="accuracy",
+        )
+
+        assert value == 0.85
+
+    def test_defaults_to_first_validation_loss(self):
+        value = _get_latest_metric_value(
+            val_losses={"mse_loss": [0.5], "l1_loss": [0.4]},
+            val_metrics={},
+        )
+
+        assert value == 0.5
+
+    def test_missing_metric_raises_error(self):
+        with pytest.raises(ValueError, match="not found in logs"):
+            _get_latest_metric_value({}, {}, "missing")
+
+    def test_missing_default_loss_raises_error(self):
+        with pytest.raises(RuntimeError, match="No validation losses"):
+            _get_latest_metric_value({}, {})
+
+    def test_empty_metric_history_raises_error(self):
+        with pytest.raises(RuntimeError, match="has no recorded values"):
+            _get_latest_metric_value({"mse_loss": []}, {}, "mse_loss")
+
+    def test_invalid_metric_history_raises_error(self):
+        with pytest.raises(TypeError, match="list or tuple"):
+            _get_latest_metric_value({"mse_loss": 0.5}, {}, "mse_loss")
+
+
+class TestEarlyStopHelper:
+    def test_improvement_resets_counter(self):
+        losses = {"val_loss": [0.5]}
+        helper = EarlyStopHelper(
+            trainer_val_losses_ref=losses,
+            best_metric_name="val_loss",
+        )
+        helper.initialize_early_stop(patience=3)
+        assert helper.update() is False
+        losses["val_loss"].append(0.6)
+        assert helper.update() is False
+        losses["val_loss"].append(0.4)
+
+        should_stop = helper.update()
+
+        assert should_stop is False
+        assert helper.counter == 0
+        assert helper.best_epoch_state.best_metric_value == 0.4
+
+    def test_non_improvement_stops_at_patience(self):
+        losses = {"val_loss": [0.3]}
+        helper = EarlyStopHelper(
+            trainer_val_losses_ref=losses,
+            best_metric_name="val_loss",
+        )
+        helper.initialize_early_stop(patience=2)
+        assert helper.update() is False
+        losses["val_loss"].append(0.4)
+        assert helper.update() is False
+        losses["val_loss"].append(0.5)
+
+        assert helper.update() is True
+        assert helper.counter == 2
+
+    def test_disabled_update_is_no_op(self, minimal_model):
+        helper = EarlyStopHelper(enabled=False)
+        helper.initialize_early_stop(patience=1)
+
+        should_stop = helper.update(epoch=1, model=minimal_model)
+
+        assert should_stop is False
+        assert helper.counter == 0
+        assert helper.best_model is None
+
+    def test_enabled_update_requires_initialization(self):
+        helper = EarlyStopHelper(trainer_val_losses_ref={"loss": [0.5]})
+
+        with pytest.raises(RuntimeError, match="has not been initialized"):
+            helper.update()
+
+    def test_initialize_rejects_invalid_patience(self):
+        helper = EarlyStopHelper()
+
+        with pytest.raises(ValueError, match="at least 1"):
+            helper.initialize_early_stop(0)
+
+    def test_reinitialize_resets_counter_and_preserves_best_state(self):
+        losses = {"loss": [0.3]}
+        helper = EarlyStopHelper(trainer_val_losses_ref=losses)
+        helper.initialize_early_stop(patience=2)
+        helper.update(epoch=1)
+        losses["loss"].append(0.4)
+        helper.update(epoch=2)
+
+        helper.initialize_early_stop(patience=3)
+
+        assert helper.counter == 0
+        assert helper.best_epoch_state.best_metric_value == 0.3
+        assert helper.best_epoch_state.best_epoch == 1

From f4f2ea4f6d5206c3da0e470d468f0799b1c209fa Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Mon, 3 Aug 2026 17:27:00 -0600
Subject: [PATCH 4/8] Update CHANGELOG for version 0.4.8: Fix best model
 updating/saving behavior and refactor early stopping logic

---
 CHANGELOG.md | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 904308a..4790da6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ---
 
+## [0.4.8] - 2026-08-03
+
+### Fixed 
+
+#### Best model updating/saving behavior (`virtual_stain_flow/trainers/`)
+- Fixes best model updates storing only a shallow copy of statedict of models as opposed to immutable statedict. 
+- Fixes best model saving happens regardless of absent validation set. 
+In previous versions if the valiadtion set is absent from training and best model cannot possibly be determined, the trainer silently updates the `.best_model` property by the most recent model weights from training and that gets saved by the logger as the "best model". 
+With this fix the `trainer.best_model` will return `None` if best model cannot be determined due to missing of validation set, and correspondingly the logger will no longer write best model artifacts if not available.  
+
+### Refactored
+
+#### Early stopping logic (`virtual_stain_flow/trainers/`)
+- The early stopping related functionalities are now moved outside of `AbstractTrainer` to standalone utility classes to facilitate unit testing and reduce size of single module. 
+
+---
+
 ## [0.4.7] - 2026-07-28
 
 ### Fixed

From aa25ce66ece217adaa56654aeac22a38142f4965 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 4 Aug 2026 17:28:52 -0600
Subject: [PATCH 5/8] Refactor early stopping logic in trainers to limit model
 cloning to only once at start when the model is moved to CPU, subsequent best
 model updates happen by copying of state dict

---
 .../trainers/AbstractTrainer.py               |   4 +-
 .../trainers/logging_gan_trainer.py           |   3 +-
 .../trainers/logging_trainer.py               |   5 +-
 .../trainers/trainer_utils/__init__.py        |   2 -
 .../trainers/trainer_utils/early_stop.py      |  82 ++++-
 .../vsf_logging/from_run.py                   | 330 ++++++++++++++++++
 tests/trainers/test_logging_trainer.py        |   8 +-
 .../trainers/trainer_utils/test_early_stop.py |  91 +++--
 8 files changed, 447 insertions(+), 78 deletions(-)
 create mode 100644 src/virtual_stain_flow/vsf_logging/from_run.py

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index a1077e5..21c89d8 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -120,6 +120,7 @@ def _init_state(
 
         # Early stopping state
         self._early_stop_helper = EarlyStopHelper(
+            model=self._model,
             best_mode=early_termination_mode,
             trainer_val_losses_ref=self._val_losses,
             trainer_val_metrics_ref=self._val_metrics,
@@ -362,7 +363,6 @@ def train(
             # Update early stopping
             should_stop = self._early_stop_helper.update(
                 epoch=self.epoch,
-                model=self.model
             )
 
             if hasattr(logger, "on_epoch_end"):
@@ -370,7 +370,7 @@ def train(
 
             if should_stop:
                 print(f"Early termination at epoch {self.epoch} "
-                      f"with best validation metric {self._early_stop_helper.best_epoch_state.best_metric_value}")
+                    f"with best validation metric {self._early_stop_helper.best_metric_value}")
                 break
 
         if hasattr(logger, "on_train_end"):
diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
index 11a6d8c..e51cc14 100644
--- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
@@ -52,7 +52,7 @@ def __init__(
         :kwargs: Additional arguments for the AbstractTrainer
         """
 
-        device = kwargs.get('device', torch.device('cpu'))
+        device = kwargs.pop('device', torch.device('cpu'))
 
         # Registry for logging model parameters
         self._models: List[torch.nn.Module] = [generator, discriminator]
@@ -85,6 +85,7 @@ def __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/logging_trainer.py b/src/virtual_stain_flow/trainers/logging_trainer.py
index 645910f..4c81ed3 100644
--- a/src/virtual_stain_flow/trainers/logging_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_trainer.py
@@ -28,7 +28,6 @@ def __init__(
         model: torch.nn.Module,
         optimizer: torch.optim.Optimizer,
         losses: Union[torch.nn.Module, List[torch.nn.Module]],
-        device: torch.device,
         loss_weights: Optional[Union[Scalar, List[Scalar]]] = None,
         **kwargs
     ):
@@ -38,10 +37,11 @@ def __init__(
         :param model: The generator model to be trained.
         :param optimizer: The optimizer to be used for training.
         :param losses: The loss function(s) to be used for training.
-        :param device: The device to run the training on.
         :param loss_weights: Optional weights for each loss function.
         :kwargs: Additional arguments for the AbstractTrainer (for data/metric and more)
         """
+
+        device = kwargs.pop('device', torch.device('cpu'))
         
         # Registry for logging model parameters
         self._models: List[torch.nn.Module] = [model]
@@ -82,6 +82,7 @@ def __init__(
         super().__init__(
             model=self._forward_group.model,
             optimizer=self._forward_group.optimizer, # type: ignore
+            device=device,
             **kwargs
         )
 
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
index d106587..f0a09fa 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
@@ -3,7 +3,6 @@
 """
 
 from .early_stop import (
-    BestEpochState,
     EarlyStopHelper,
     _get_latest_metric_value,
 )
@@ -11,7 +10,6 @@
 
 
 __all__ = [
-    "BestEpochState",
     "EarlyStopHelper",
     "_get_latest_metric_value",
     "save_model",
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py b/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
index a2a6841..d2bc706 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py
@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Literal, Optional
 
 import torch
@@ -7,23 +7,19 @@
 
 
 @dataclass
-class BestEpochState:
+class _BestEpochState:
     """Track the best validation observation and its model snapshot."""
 
+    model_copy: Optional[torch.nn.Module]
     best_mode: Literal["min", "max"] = "min"
     best_metric_value: Optional[float] = None
     best_epoch: int = 0
-    best_model: Optional[torch.nn.Module] = field(default=None, init=False)
+    enabled: bool = True
 
     def __post_init__(self) -> None:
         if self.best_mode not in ("min", "max"):
             raise ValueError("best_mode must be either 'min' or 'max'.")
 
-        if self.best_metric_value is None:
-            self.best_metric_value = (
-                float("inf") if self.best_mode == "min" else float("-inf")
-            )
-
     def update(
         self,
         metric_value: float,
@@ -32,34 +28,70 @@ def update(
     ) -> bool:
         """Record an improved observation and return whether it improved."""
         metric_value = float(metric_value)
-        is_better = (
-            self.best_mode == "min" and metric_value < self.best_metric_value
-        ) or (
-            self.best_mode == "max" and metric_value > self.best_metric_value
-        )
+        is_better = self._is_better(metric_value)
         if not is_better:
             return False
 
         self.best_metric_value = metric_value
         self.best_epoch = epoch if epoch is not None else self.best_epoch + 1
-        if model is not None:
-            self.best_model = clone_model(model)
+        if self.model_copy is not None:
+            self._update_best_model_state(model)
 
         return True
 
+    def _is_better(self, metric_value: float) -> bool:
+        """Return whether the incoming metric improves the current best value."""
+        if self.best_metric_value is None:
+            return True
+        if self.best_mode == "min":
+            return metric_value < self.best_metric_value
+        return metric_value > self.best_metric_value
+
+    @property
+    def best_model(self) -> Optional[torch.nn.Module]:
+        """Return the best model snapshot, if an observation was recorded."""
+        if not self.enabled or self.best_metric_value is None or self.model_copy is None:
+            return None
+        return self.model_copy
+
+    @torch.no_grad()
+    def _update_best_model_state(
+        self,
+        model: torch.nn.Module,
+    ) -> None:
+        if model is None:
+            raise ValueError("A valid PyTorch model reference is required for snapshotting.")
+        self.model_copy.load_state_dict(
+            model.state_dict(),
+            strict=True,
+            assign=False,
+        )
+
 
 class EarlyStopHelper:
     """Track validation improvements and decide when training should stop."""
 
     def __init__(
         self,
+        model: torch.nn.Module,
         best_mode: Literal["min", "max"] = "min",
         trainer_val_losses_ref: Optional[dict] = None,
         trainer_val_metrics_ref: Optional[dict] = None,
         best_metric_name: Optional[str] = None,
         enabled: bool = True,
     ) -> None:
-        self.best_epoch_state = BestEpochState(best_mode=best_mode)
+        if model is None or not isinstance(model, torch.nn.Module):
+            raise ValueError("A valid PyTorch model must be provided.")
+
+        # Store model reference and allocate a snapshot once when enabled.
+        self.model_reference = model
+        model_copy = self._create_cpu_snapshot(self.model_reference) if enabled else None
+
+        self.best_epoch_state = _BestEpochState(
+            model_copy=model_copy,
+            best_mode=best_mode,
+            enabled=enabled,
+        )
         self.trainer_val_losses_ref = (
             trainer_val_losses_ref if trainer_val_losses_ref is not None else {}
         )
@@ -82,7 +114,6 @@ def initialize_early_stop(self, patience: int) -> None:
     def update(
         self,
         epoch: Optional[int] = None,
-        model: Optional[torch.nn.Module] = None,
     ) -> bool:
         """Update tracking state and return whether training should stop."""
         if not self.enabled:
@@ -98,13 +129,23 @@ def update(
             self.trainer_val_metrics_ref,
             self.best_metric_name,
         )
-        if self.best_epoch_state.update(metric_value, epoch, model):
+        if self.best_epoch_state.update(metric_value, epoch, self.model_reference):
             self._counter = 0
         else:
             self._counter += 1
 
         return self._counter >= self._patience
 
+    @staticmethod
+    def _create_cpu_snapshot(model: torch.nn.Module) -> torch.nn.Module:
+        """Create a CPU snapshot while restoring the original model device."""
+        first_param = next(model.parameters(), None)
+        model_device = first_param.device if first_param is not None else torch.device("cpu")
+        model.to("cpu")
+        model_copy = clone_model(model)
+        model.to(model_device)
+        return model_copy
+
     @property
     def best_model(self) -> Optional[torch.nn.Module]:
         """Return the best model snapshot, if an observation was recorded."""
@@ -115,6 +156,11 @@ def counter(self) -> int:
         """Return the number of consecutive observations without improvement."""
         return self._counter
 
+    @property
+    def best_metric_value(self) -> Optional[float]:
+        """Return the best tracked metric value, if any."""
+        return self.best_epoch_state.best_metric_value
+
 
 def _get_latest_metric_value(
     val_losses: dict,
diff --git a/src/virtual_stain_flow/vsf_logging/from_run.py b/src/virtual_stain_flow/vsf_logging/from_run.py
new file mode 100644
index 0000000..52d17a5
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/from_run.py
@@ -0,0 +1,330 @@
+"""
+Helper module for re-loading model from MLflow tracking info
+"""
+
+import pathlib
+import re
+import json
+import importlib
+from typing import Optional, Literal, Dict
+
+import torch
+from mlflow.tracking import MlflowClient
+from mlflow.entities.file_info import FileInfo
+
+from ..evaluation.as_gif import images_to_numbered_gif
+
+
+def _artifact_sort_key(artifact_path: str):
+    p = pathlib.Path(artifact_path)
+    nums = [int(x) for x in re.findall(r"\d+", p.stem)]
+    return (nums, p.stem)
+
+
+def _get_weight_artifacts(
+    client: MlflowClient,
+    tracking_run_id: str
+) -> list[FileInfo]:
+    
+    try:
+        weight_artifacts = [
+            item for item in client.list_artifacts(tracking_run_id, path="weights")
+            if (not item.is_dir) and pathlib.Path(item.path).suffix.lower() in {".pt", ".pth", ".ckpt", ".bin"}
+        ]
+    except Exception as e:
+        raise ValueError(f"Failed to list artifacts for run ID '{tracking_run_id}': {e}")
+    
+    if not weight_artifacts:
+        raise ValueError(f"No weight artifacts found for run ID '{tracking_run_id}' at path 'weights/'")
+
+    return weight_artifacts
+
+
+def _select_weight_artifact(
+    client: MlflowClient, 
+    tracking_run_id: str,
+    load_weight_mode: Literal["latest", "best"] = "latest"
+) -> FileInfo:
+    
+    weight_artifacts = _get_weight_artifacts(client, tracking_run_id)
+
+    if load_weight_mode == "best":
+        best_weight_artifact = [
+            weight_file
+            for weight_file in weight_artifacts
+            if "best" in weight_file.path.lower()
+        ]
+        if not best_weight_artifact:
+            raise ValueError(
+                "No weight artifacts with 'best' in filename found "
+                f"for run ID '{tracking_run_id}'")
+        if len(best_weight_artifact) > 1:
+            raise RuntimeError(
+                "Multiple weight artifacts with 'best' in filename found "
+                f"for run ID '{tracking_run_id}': {[a.path for a in best_weight_artifact]}")
+        
+        return best_weight_artifact[0]
+    
+    elif load_weight_mode == "latest":
+
+        latest_weight_artifact = sorted(weight_artifacts, key=lambda x: _artifact_sort_key(x.path))[-1]
+
+        return latest_weight_artifact
+    
+
+def _load_config(
+    client: MlflowClient,
+    tracking_run_id: str
+) -> Dict:
+    
+    try:
+        config_artifacts = client.list_artifacts(tracking_run_id, path="configs")
+    except Exception as e:
+        raise ValueError(f"Failed to list config artifacts for run ID '{tracking_run_id}': {e}")
+    
+    if not config_artifacts:
+        raise ValueError(f"No config artifacts found for run ID '{tracking_run_id}' at path 'configs/'")
+
+    generator_config_artifacts = [
+        artifact for artifact in config_artifacts
+        if (artifact.path.lower().endswith(".json") and all(keyword not in artifact.path.lower() for keyword in ["discriminator", "loss_group", "optimizer"]))
+    ]
+    if not generator_config_artifacts:
+        raise ValueError(f"No generator config artifacts found for run ID '{tracking_run_id}'")
+    if len(generator_config_artifacts) > 1:
+        raise ValueError(f"Multiple generator config artifacts found for run ID '{tracking_run_id}': {[a.path for a in generator_config_artifacts]}")
+
+    # TODO: need to extend support for multiple config files
+    # Best place to start is probably not here but in the MLflow logging process
+    # to ensure more consistent naming of config artifacts
+    try:
+        local_config_path = client.download_artifacts(tracking_run_id, generator_config_artifacts[0].path)
+    except Exception as e:
+        raise ValueError(f"Failed to download config artifact '{generator_config_artifacts[0].path}' for run ID '{tracking_run_id}': {e}")
+    
+    try:
+        config = json.load(open(local_config_path, "r"))
+    except Exception as e:
+        raise ValueError(f"Failed to load config from '{local_config_path}': {e}")
+    
+    return config
+
+
+def _artifact_path_exists(
+    client: MlflowClient,
+    tracking_run_id: str,
+    artifact_path: str,
+) -> bool:
+    try:
+        items = client.list_artifacts(tracking_run_id, path=artifact_path)
+    except Exception:
+        return False
+
+    return len(items) > 0
+
+
+# maybe this helper belong better somewhere else
+def _get_class_from_path(class_path: str):
+    module_path, class_name = class_path.rsplit(".", 1)
+
+    try:
+        module = importlib.import_module(module_path)
+        cls = getattr(module, class_name)
+    except Exception as e:
+        raise ValueError(f"Failed to get class from path '{class_path}': {e}")    
+    
+    return cls
+
+
+def from_mlflow_run(
+    tracking_run_id: str,
+    tracking_uri: str,
+    experiment_name: str,
+    model_class: Optional[torch.nn.Module] = None,
+    model_kwargs: Optional[Dict] = None,
+    load_weight_mode: Literal["latest", "best"] = "latest",
+    device: torch.device = 'cpu'
+) -> torch.nn.Module:
+    
+    try:
+        client = MlflowClient(tracking_uri=tracking_uri)
+    except Exception as e:
+        raise ValueError(f"Failed to initialize MLflow client: {e}")
+    
+    try:
+        experiment = client.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            raise ValueError(f"Experiment '{experiment_name}' not found at URI '{tracking_uri}'")
+    except Exception as e:
+        raise ValueError(f"Failed to retrieve experiment '{experiment_name}': {e}")
+    
+    try:
+        run = client.get_run(tracking_run_id)
+        if run.info.experiment_id != experiment.experiment_id:
+            raise ValueError(
+                f"Run {tracking_run_id} belongs to experiment_id={run.info.experiment_id}, "
+                f"not '{experiment_name}' (id={experiment.experiment_id})."
+            )
+    except Exception as e:
+        raise ValueError(f"Failed to retrieve run with ID '{tracking_run_id}': {e}")
+    
+    selected_artifact = _select_weight_artifact(client, tracking_run_id, load_weight_mode)
+
+    try:
+        local_weight_path = client.download_artifacts(tracking_run_id, selected_artifact.path)
+    except Exception as e:
+        raise ValueError(f"Failed to download artifact '{selected_artifact.path}' for run ID '{tracking_run_id}': {e}")
+    
+    try:
+        ckpt = torch.load(local_weight_path, map_location=device)
+    except Exception as e:
+        raise ValueError(f"Failed to load checkpoint from '{local_weight_path}': {e}")
+    
+    if isinstance(ckpt, dict):
+        if "generator_state_dict" in ckpt:
+            state_dict = ckpt["generator_state_dict"]
+        elif "model_state_dict" in ckpt:
+            state_dict = ckpt["model_state_dict"]
+        elif "state_dict" in ckpt:
+            state_dict = ckpt["state_dict"]
+        else:
+            # assume raw state_dict-like checkpoint
+            state_dict = ckpt
+    else:
+        raise TypeError(f"Unsupported checkpoint type: {type(ckpt)}")
+    
+    if model_class is not None and model_kwargs is not None:
+        # case 1: user supplies both model class and kwargs
+        # override config class path and kwargs if explicitly provided
+        try:
+            model = model_class(**model_kwargs)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to instantiate model from class '{model_class}': {e} "
+                f"with custom model specification {model_class} and {model_kwargs}"
+            )
+    elif model_class is not None and model_kwargs is None:
+        # case 2: user supplies model class only
+        # infer kwargs from config best effort and let it crash if needed
+        # kwargs cannot be inferred
+        config = _load_config(client, tracking_run_id)
+        try:
+            model = model_class(**config.get("init", {}))
+        except Exception as e:
+            raise ValueError(
+                f"Failed to instantiate model from class '{model_class}': {e} "
+                f"with config specification {config.get('init', {})}"
+            )
+    elif model_class is None and model_kwargs is not None:
+        # case 3: user supplies model kwargs without class
+        # attempt to infer class from config but use override kwargs
+        try:
+            config = _load_config(client, tracking_run_id)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to load config for run ID '{tracking_run_id}': {e} "
+                "for model class inferring"
+            )
+        
+        try:
+            model_class = _get_class_from_path(config['class_path'])
+        except Exception as e:
+            raise ValueError(
+                f"Failed to get model class from config for run ID '{tracking_run_id}': {e} "
+                "for model class inferring"
+            )
+        
+        try:
+            model = model_class(**model_kwargs)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to instantiate model from class '{model_class}': {e} "
+                f"with override kwargs {model_kwargs}"
+            )
+    else:
+        # case 4: user does not supply model class or kwargs
+        # attempt to infer both from config but let it crash if needed
+        try:
+            config = _load_config(client, tracking_run_id)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to load config for run ID '{tracking_run_id}': {e} "
+                "for model class and kwargs inferring"
+            )
+        
+        try:
+            model_class = _get_class_from_path(config['class_path'])
+        except Exception as e:
+            raise ValueError(
+                f"Failed to get model class from config for run ID '{tracking_run_id}': {e} "
+                "for model class inferring"
+            )
+        
+        try:
+            model = model_class(**config.get("init", {}))
+        except Exception as e:
+            raise ValueError(
+                f"Failed to instantiate model from class '{model_class}': {e} "
+                f"with config specification {config.get('init', {})}"
+            )
+        
+    try:
+        load_info = model.load_state_dict(state_dict)
+    except Exception as e:
+        raise ValueError(f"Failed to load state dict into model: {e}")
+    
+    return model, load_info
+
+
+def gifs_from_mlflow_run(
+    tracking_run_id: str,
+    tracking_uri: str,
+    output_dir: pathlib.Path,
+    pattern: str = "epoch",
+    fps: float = 5,
+    number_color: str = "black",
+    subset: Optional[list[int]] = None,
+    font_size: int = 24,
+    padding: int = 8,
+    loop: int = 0,
+) -> dict[str, pathlib.Path]:
+    """
+    Generate GIFs from common prediction plot folders in MLflow artifacts.
+    """
+
+    client = MlflowClient(tracking_uri=tracking_uri)
+    output_dir = pathlib.Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    artifact_paths = [
+        "plots/epoch/plot_train_predictions",
+        "plots/epoch/plot_heldout_predictions",
+    ]
+
+    outputs: dict[str, pathlib.Path] = {}
+
+    for artifact_path in artifact_paths:
+        if not _artifact_path_exists(client, tracking_run_id, artifact_path):
+            continue
+
+        local_dir = pathlib.Path(
+            client.download_artifacts(tracking_run_id, artifact_path)
+        )
+
+        output_path = output_dir / f"{tracking_run_id}_{local_dir.name}.gif"
+
+        images_to_numbered_gif(
+            image_dir=local_dir,
+            output_path=output_path,
+            pattern=pattern,
+            fps=fps,
+            number_color=number_color,
+            subset=subset,
+            font_size=font_size,
+            padding=padding,
+            loop=loop,
+        )
+
+        outputs[artifact_path] = output_path
+
+    return outputs
diff --git a/tests/trainers/test_logging_trainer.py b/tests/trainers/test_logging_trainer.py
index ff1f44d..c4c455b 100644
--- a/tests/trainers/test_logging_trainer.py
+++ b/tests/trainers/test_logging_trainer.py
@@ -212,9 +212,9 @@ def test_save_current_model_creates_file(
             assert len(paths) == 1
             assert paths[0].exists()
 
-    def test_save_best_model_without_snapshot_returns_empty_list(
+    def test_save_best_model_without_validation_returns_empty_list(
         self, mock_model_with_save, mock_optimizer, simple_loss,
-        train_dataloader, val_dataloader
+        train_dataloader,
     ):
         """Test that saving a best model requires an established best snapshot."""
         from virtual_stain_flow.trainers.logging_trainer import SingleGeneratorTrainer
@@ -225,7 +225,7 @@ def test_save_best_model_without_snapshot_returns_empty_list(
             losses=simple_loss,
             device=torch.device('cpu'),
             train_loader=train_dataloader,
-            val_loader=val_dataloader,
+            val_loader=None,
             batch_size=2
         )
 
@@ -255,7 +255,7 @@ def test_save_best_model_creates_file_after_best_is_established(
         )
         trainer._early_stop_helper.initialize_early_stop(patience=1)
         trainer.val_losses["MSELoss"].append(0.5)
-        trainer._early_stop_helper.update(epoch=1, model=trainer.model)
+        trainer._early_stop_helper.update(epoch=1)
 
         with tempfile.TemporaryDirectory() as tmpdir:
             paths = trainer.save_model(
diff --git a/tests/trainers/trainer_utils/test_early_stop.py b/tests/trainers/trainer_utils/test_early_stop.py
index 941f848..e62d7f0 100644
--- a/tests/trainers/trainer_utils/test_early_stop.py
+++ b/tests/trainers/trainer_utils/test_early_stop.py
@@ -2,53 +2,11 @@
 import torch
 
 from virtual_stain_flow.trainers.trainer_utils.early_stop import (
-    BestEpochState,
     EarlyStopHelper,
     _get_latest_metric_value,
 )
 
 
-class TestBestEpochState:
-    @pytest.mark.parametrize(
-        ("best_mode", "initial_value", "new_value"),
-        [("min", 0.5, 0.3), ("max", 0.8, 0.9)],
-    )
-    def test_update_records_improvement(
-        self, best_mode, initial_value, new_value
-    ):
-        state = BestEpochState(
-            best_mode=best_mode,
-            best_metric_value=initial_value,
-        )
-
-        is_better = state.update(new_value, epoch=4)
-
-        assert is_better is True
-        assert state.best_metric_value == new_value
-        assert state.best_epoch == 4
-
-    def test_update_ignores_worse_value(self):
-        state = BestEpochState(best_metric_value=0.3, best_epoch=2)
-
-        is_better = state.update(0.5, epoch=3)
-
-        assert is_better is False
-        assert state.best_metric_value == 0.3
-        assert state.best_epoch == 2
-
-    def test_update_snapshots_model(self, minimal_model):
-        state = BestEpochState()
-
-        state.update(0.5, model=minimal_model)
-
-        assert state.best_model is not None
-        assert state.best_model is not minimal_model
-        for copied_parameter, parameter in zip(
-            state.best_model.parameters(), minimal_model.parameters()
-        ):
-            assert torch.equal(copied_parameter, parameter)
-
-
 class TestGetLatestMetricValue:
     def test_gets_value_from_losses(self):
         value = _get_latest_metric_value(
@@ -94,9 +52,21 @@ def test_invalid_metric_history_raises_error(self):
 
 
 class TestEarlyStopHelper:
+    def test_best_model_is_none_before_first_improvement(self):
+        losses = {"val_loss": [0.5]}
+        helper = EarlyStopHelper(
+            model=torch.nn.Linear(1, 1),
+            trainer_val_losses_ref=losses,
+            best_metric_name="val_loss",
+        )
+        helper.initialize_early_stop(patience=2)
+
+        assert helper.best_model is None
+
     def test_improvement_resets_counter(self):
         losses = {"val_loss": [0.5]}
         helper = EarlyStopHelper(
+            model=torch.nn.Linear(1, 1),
             trainer_val_losses_ref=losses,
             best_metric_name="val_loss",
         )
@@ -110,11 +80,12 @@ def test_improvement_resets_counter(self):
 
         assert should_stop is False
         assert helper.counter == 0
-        assert helper.best_epoch_state.best_metric_value == 0.4
+        assert helper.best_metric_value == 0.4
 
     def test_non_improvement_stops_at_patience(self):
         losses = {"val_loss": [0.3]}
         helper = EarlyStopHelper(
+            model=torch.nn.Linear(1, 1),
             trainer_val_losses_ref=losses,
             best_metric_name="val_loss",
         )
@@ -128,30 +99,52 @@ def test_non_improvement_stops_at_patience(self):
         assert helper.counter == 2
 
     def test_disabled_update_is_no_op(self, minimal_model):
-        helper = EarlyStopHelper(enabled=False)
+        helper = EarlyStopHelper(model=minimal_model, enabled=False)
         helper.initialize_early_stop(patience=1)
 
-        should_stop = helper.update(epoch=1, model=minimal_model)
+        should_stop = helper.update(epoch=1)
 
         assert should_stop is False
         assert helper.counter == 0
         assert helper.best_model is None
 
+    def test_snapshot_tracks_internal_model_reference(self, minimal_model):
+        losses = {"val_loss": [0.5]}
+        helper = EarlyStopHelper(
+            model=minimal_model,
+            trainer_val_losses_ref=losses,
+            best_metric_name="val_loss",
+        )
+        helper.initialize_early_stop(patience=2)
+
+        with torch.no_grad():
+            for parameter in minimal_model.parameters():
+                parameter.add_(1.0)
+
+        helper.update(epoch=1)
+
+        assert helper.best_model is not None
+        assert helper.best_model is not minimal_model
+        for copied_parameter, parameter in zip(
+            helper.best_model.parameters(), minimal_model.parameters()
+        ):
+            assert torch.equal(copied_parameter, parameter)
+
     def test_enabled_update_requires_initialization(self):
-        helper = EarlyStopHelper(trainer_val_losses_ref={"loss": [0.5]})
+        helper = EarlyStopHelper(model=torch.nn.Linear(1, 1),trainer_val_losses_ref={"loss": [0.5]})
 
         with pytest.raises(RuntimeError, match="has not been initialized"):
             helper.update()
 
     def test_initialize_rejects_invalid_patience(self):
-        helper = EarlyStopHelper()
+        helper = EarlyStopHelper(model=torch.nn.Linear(1, 1))
 
         with pytest.raises(ValueError, match="at least 1"):
             helper.initialize_early_stop(0)
 
     def test_reinitialize_resets_counter_and_preserves_best_state(self):
         losses = {"loss": [0.3]}
-        helper = EarlyStopHelper(trainer_val_losses_ref=losses)
+        helper = EarlyStopHelper(model=torch.nn.Linear(1, 1), trainer_val_losses_ref=losses)
         helper.initialize_early_stop(patience=2)
         helper.update(epoch=1)
         losses["loss"].append(0.4)
@@ -160,5 +153,5 @@ def test_reinitialize_resets_counter_and_preserves_best_state(self):
         helper.initialize_early_stop(patience=3)
 
         assert helper.counter == 0
-        assert helper.best_epoch_state.best_metric_value == 0.3
+        assert helper.best_metric_value == 0.3
         assert helper.best_epoch_state.best_epoch == 1

From f7b9248ac7a823fbb4dc231c334c1c0a52530ee4 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 4 Aug 2026 17:35:48 -0600
Subject: [PATCH 6/8] Enhance early termination logic to allow best model
 updates without a validation set; add validation presence check for early
 termination metric.

---
 src/virtual_stain_flow/trainers/AbstractTrainer.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 21c89d8..932faf3 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -72,6 +72,11 @@ def __init__(
             early-termination count on the validation dataset. 
             If None, early termination is disabled and the
             training will run for the specified number of epochs.
+            This metric also controls best model updating which will
+            be reflected in the best_model property and the save_model method.
+            However, unlike the early termination, the best model will be
+            updated even if early_termination_metric is None, 
+            so that the best model can be saved at the end of training.
         :param early_termination_mode: (optional)
         """
 
@@ -118,6 +123,13 @@ def _init_state(
         self._train_metrics = defaultdict(list)
         self._val_metrics = defaultdict(list)
 
+        validation_present = bool(self._val_loader)
+        if early_termination_metric is not None and not validation_present:
+            raise RuntimeError(
+                "Cannot specify early_termination_metric if validation set or loader is not supplied. "
+                "Please either provide a validation dataset/loader or set early_termination_metric to None."
+            )
+
         # Early stopping state
         self._early_stop_helper = EarlyStopHelper(
             model=self._model,
@@ -125,7 +137,7 @@ def _init_state(
             trainer_val_losses_ref=self._val_losses,
             trainer_val_metrics_ref=self._val_metrics,
             best_metric_name=early_termination_metric,
-            enabled=bool(self._val_loader),
+            enabled=validation_present,
         )
 
         return None

From a9729593bab076db22f9194fd69c56e9ef86ee59 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 4 Aug 2026 17:40:47 -0600
Subject: [PATCH 7/8] Update pre-commit workflow to install dependencies from
 dev requirements and add CPU-safe test execution

---
 .github/workflows/pre-commit.yml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index cd1f5b0..a303974 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -20,7 +20,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pre-commit
+          pip install -e ".[dev]"
 
       - name: Run pre-commit
         uses: pre-commit/action@v3.0.1
@@ -28,3 +28,6 @@ jobs:
       - name: pre-commit-ci-lite
         uses: pre-commit-ci/lite-action@v1.1.0
         if: always()
+
+      - name: Run CPU-safe tests
+        run: pytest -m "not gpu"

From 059619d24693d55c8f41391f0d3cbb5e7411ad0a Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 4 Aug 2026 17:43:47 -0600
Subject: [PATCH 8/8] Refactor pre-commit workflow to install pre-commit
 dependencies and separate CPU-safe test execution steps

---
 .github/workflows/pre-commit.yml | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index a303974..e29d2a8 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -20,7 +20,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -e ".[dev]"
+          pip install pre-commit
 
       - name: Run pre-commit
         uses: pre-commit/action@v3.0.1
@@ -29,5 +29,22 @@ jobs:
         uses: pre-commit-ci/lite-action@v1.1.0
         if: always()
 
+  tests-cpu:
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install package + test deps
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+
       - name: Run CPU-safe tests
         run: pytest -m "not gpu"