diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index cd1f5b0..e29d2a8 100644 --- a/.github/workflows/pre-commit.yml +++ b/.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" 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 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 diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py index 36c25ae..932faf3 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, ): @@ -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 diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py index d959457..e51cc14 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 @@ -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): """ diff --git a/src/virtual_stain_flow/trainers/logging_trainer.py b/src/virtual_stain_flow/trainers/logging_trainer.py index ff20c8d..4c81ed3 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 @@ -29,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 ): @@ -39,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] @@ -83,6 +82,7 @@ def __init__( super().__init__( model=self._forward_group.model, optimizer=self._forward_group.optimizer, # type: ignore + device=device, **kwargs ) @@ -151,27 +151,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/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..f0a09fa --- /dev/null +++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py @@ -0,0 +1,16 @@ +""" +Trainer utilities for virtual stain flow. +""" + +from .early_stop import ( + EarlyStopHelper, + _get_latest_metric_value, +) +from .save_model import save_model + + +__all__ = [ + "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..d2bc706 --- /dev/null +++ b/src/virtual_stain_flow/trainers/trainer_utils/early_stop.py @@ -0,0 +1,193 @@ +from dataclasses import dataclass +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.""" + + model_copy: Optional[torch.nn.Module] + best_mode: Literal["min", "max"] = "min" + best_metric_value: Optional[float] = None + best_epoch: int = 0 + 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'.") + + 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._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 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: + 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 {} + ) + 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, + ) -> 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, 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.""" + return self.best_epoch_state.best_model + + @property + 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, + 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] 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_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..c4c455b 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_validation_returns_empty_list( + self, mock_model_with_save, mock_optimizer, simple_loss, + train_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=None, + 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) + + 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..e62d7f0 --- /dev/null +++ b/tests/trainers/trainer_utils/test_early_stop.py @@ -0,0 +1,157 @@ +import pytest +import torch + +from virtual_stain_flow.trainers.trainer_utils.early_stop import ( + EarlyStopHelper, + _get_latest_metric_value, +) + + +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_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", + ) + 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_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", + ) + 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(model=minimal_model, enabled=False) + helper.initialize_early_stop(patience=1) + + 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(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(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(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) + helper.update(epoch=2) + + helper.initialize_early_stop(patience=3) + + assert helper.counter == 0 + assert helper.best_metric_value == 0.3 + assert helper.best_epoch_state.best_epoch == 1