Fix best model

.github/workflows/pre-commit.yml

@@ -28,3 +28,23 @@ jobs:
       - name: pre-commit-ci-lite
         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"

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

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

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,
     ):
@@ -71,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)
         """
 
@@ -100,17 +106,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 +123,23 @@ 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,
+            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=validation_present,
+        )
+
         return None
 
     def _init_data(
@@ -308,11 +324,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 +372,21 @@ def train(
                     f"val_{metric_name}", val_metric, epoch
                 )
 
+            # Update early stopping
+            should_stop = self._early_stop_helper.update(
+                epoch=self.epoch,
+            )
+
             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_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 +403,6 @@ def _update_epoch_progress(
             f"{phase} Batch {batch_idx + 1}/{num_batches}"
         )
 
-    @abstractmethod
     def save_model(
         self,
         save_path: pathlib.Path,
@@ -439,7 +411,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 +479,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 +515,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

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
@@ -53,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]
@@ -86,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
         )
 
@@ -200,27 +200,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):
     """