diff --git a/CHANGELOG.md b/CHANGELOG.md index 4790da6..1de026b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,21 @@ All notable chagnes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.9] - 2026-09-11 + +### Added + + +#### Initial training resume support (`virtual_stain_flow/vsf_logging/`, `virtual_stain_flow/trainers/`) +- Added MLflow logging for optimizer state to support resuming training runs from logged artifacts and trainer state. + +### Fixed + +#### Pandas validation API update (`virtual_stain_flow/datasets/ds_engine/`) +- Replaced deprecated `DataFrame.applymap` usage with `DataFrame.map` in dataset input validation to keep the GH Actions test suite green. + + + --- ## [0.4.8] - 2026-08-03 diff --git a/pyproject.toml b/pyproject.toml index 56bd4fe..a9497a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,13 +6,13 @@ build-backend = "setuptools.build_meta" [project] name = "virtual_stain_flow" -version = "0.4.7" +version = "0.4.9" description = "For developing virtual staining models" requires-python = ">=3.9" dependencies = [ "pyarrow", "numpy", - "pandas", + "pandas>=2.1", "seaborn", "matplotlib", "scikit-learn", diff --git a/src/virtual_stain_flow/datasets/ds_engine/input_validation.py b/src/virtual_stain_flow/datasets/ds_engine/input_validation.py index b55c38a..9bfaf62 100644 --- a/src/virtual_stain_flow/datasets/ds_engine/input_validation.py +++ b/src/virtual_stain_flow/datasets/ds_engine/input_validation.py @@ -16,7 +16,7 @@ def _cell_contains_pathlike(x: Any, *, check_exists: bool) -> bool: """ Validate that a cell contains a valid path-like object. - Intended to be used as a lambda function in the DataFrame.applymap method. + Intended to be used as a lambda function in the DataFrame.map method. :param x: The cell value to check. :param check_exists: Whether to check if the path exists. diff --git a/src/virtual_stain_flow/datasets/ds_engine/manifest.py b/src/virtual_stain_flow/datasets/ds_engine/manifest.py index 39b3884..36bb81a 100644 --- a/src/virtual_stain_flow/datasets/ds_engine/manifest.py +++ b/src/virtual_stain_flow/datasets/ds_engine/manifest.py @@ -179,12 +179,12 @@ def get_image_dimensions( def _serialize_file_index(self) -> pd.DataFrame: """Serialize file_index to pd.DataFrame""" - return self.file_index.copy().applymap(lambda x: str(x)) + return self.file_index.map(str) @staticmethod def _deserialize_file_index(file_index: pd.DataFrame) -> pd.DataFrame: """Deserialize file_index from pd.DataFrame""" - return file_index.applymap(lambda x: Path(x)) + return file_index.map(Path) def to_config(self) -> Dict[str, Any]: """Serialize to dict""" diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py index 932faf3..f45741a 100644 --- a/src/virtual_stain_flow/trainers/AbstractTrainer.py +++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py @@ -13,7 +13,7 @@ from torch.utils.data import DataLoader from .trainer_protocol import TrainerProtocol -from .trainer_utils import EarlyStopHelper, save_model +from .trainer_utils import EarlyStopHelper, save_model, save_optimizer_state from ..metrics.AbstractMetrics import AbstractMetrics from ..engine.progress import Progress from ..datasets.data_split import default_random_split @@ -42,6 +42,7 @@ def __init__( test_ratio: Optional[float] = 0.15, metrics: Dict[str, AbstractMetrics] = None, device: Optional[torch.device] = None, + epoch: int = 0, early_termination_metric: Optional[str] = None, early_termination_mode: Literal['min', 'max'] = "min", **kwargs, @@ -68,6 +69,8 @@ def __init__( dataset is provided. Default is 0.15. :param metrics: Dictionary of metrics to be logged. :param device: (optional) The device to be used for training. + :param epoch: (optional) The starting epoch for training. + Useful for resuming training from a checkpoint. :param early_termination_metric: (optional) The metric to update early-termination count on the validation dataset. If None, early termination is disabled and the @@ -102,20 +105,25 @@ def __init__( **kwargs ) self._init_state( + epoch, early_termination_metric, early_termination_mode, **kwargs) def _init_state( self, + epoch: int, early_termination_metric: Optional[str] = None, early_termination_mode: Literal['min', 'max'] = "min", **kwargs ): + if epoch is None: + raise TypeError("epoch must be an integer, not None.") + # Epoch state - self._epoch = 0 + self._epoch = epoch # Progress tracking for loss weight scheduling - self._progress = Progress(epoch=0, step=0) + self._progress = Progress(epoch=epoch, step=0) # Loss and metrics state self._train_losses = defaultdict(list) @@ -183,7 +191,6 @@ def _init_data( **kwargs ) - self._batch_size = batch_size self._train_ratio, self._val_ratio, self._test_ratio = ( train_ratio, val_ratio, test_ratio ) @@ -194,8 +201,24 @@ def _init_data( "or provide at least train_loader." ) + self._batch_size = self._train_loader.batch_size if hasattr(self._train_loader, 'batch_size') else None + self._train_n = self._get_dataset_size(self._train_loader) + self._val_n = self._get_dataset_size(self._val_loader) + self._test_n = self._get_dataset_size(self._test_loader) + return None + @staticmethod + def _get_dataset_size(loader) -> Optional[int]: + dataset = getattr(loader, 'dataset', None) + if dataset is None: + return None + + try: + return len(dataset) + except TypeError: + return None + @abstractmethod def train_step(self, inputs: torch.Tensor, targets: torch.Tensor)->Dict[str, float]: """ @@ -323,17 +346,17 @@ def train( if hasattr(logger, "on_train_start"): logger.on_train_start() - self._epochs = epochs + epoch_range = range(self.epoch + 1, self.epoch + epochs + 1) self._epoch_pbar: Optional[tqdm] = tqdm( - range(epochs), desc="Training", unit="epoch") if verbose else None - iterable = self._epoch_pbar if self._epoch_pbar else range(epochs) + epoch_range, desc="Training", unit="epoch") if verbose else None + iterable = self._epoch_pbar if self._epoch_pbar else epoch_range self._early_stop_helper.initialize_early_stop(patience=patience if patience else epochs) for epoch in iterable: - # Increment the epoch counter - self.epoch += 1 + # Synchronize trainer state for loggers, callbacks, and schedulers + self.epoch = epoch # Invoke the on_epoch_start method of the logge if hasattr(logger, "on_epoch_start"): @@ -406,7 +429,7 @@ def _update_epoch_progress( def save_model( self, save_path: pathlib.Path, - file_name_prefix: Optional[str] = None, + file_name_prefix: str = 'generator', file_name_suffix: Optional[str] = None, file_ext: str = '.pth', best_model: bool = True @@ -414,12 +437,36 @@ def save_model( return save_model( self, save_path=save_path, - file_name_prefix=file_name_prefix or 'generator', + file_name_prefix=file_name_prefix, file_name_suffix=file_name_suffix, file_ext=file_ext, save_best_model=best_model ) + def save_optimizer_state( + self, + save_path: pathlib.Path, + file_name_prefix: str = 'optimizer', + file_name_suffix: Optional[str] = None, + file_ext: str = '.pth', + recent: bool = True + ) -> Optional[List[pathlib.Path]]: + """ + Save the optimizer state to the specified path. + """ + if not recent: + raise NotImplementedError( + "Saving non-recent optimizer states is not implemented yet." + ) + file_name_suffix = file_name_suffix or 'recent' + return save_optimizer_state( + trainer=self, + save_path=save_path, + file_name_prefix=file_name_prefix, + file_name_suffix=file_name_suffix, + file_ext=file_ext + ) + """ Log property """ @@ -452,6 +499,18 @@ def val_ratio(self): @property def test_ratio(self): return self._test_ratio + + @property + def train_n(self): + return self._train_n + + @property + def val_n(self): + return self._val_n + + @property + def test_n(self): + return self._test_n @property def model(self): @@ -468,11 +527,7 @@ def device(self): @property def batch_size(self): return self._batch_size - - @property - def epochs(self): - return self._epochs - + @property def patience(self): return self._patience diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py index e51cc14..ca9a25e 100644 --- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py +++ b/src/virtual_stain_flow/trainers/logging_gan_trainer.py @@ -84,7 +84,6 @@ def __init__( super().__init__( model=generator, # register generator as main model for early stopping optimizer=generator_optimizer, - losses=generator_loss_group, device=device, **kwargs ) diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py index 5fc3cb0..65a5143 100644 --- a/src/virtual_stain_flow/trainers/trainer_protocol.py +++ b/src/virtual_stain_flow/trainers/trainer_protocol.py @@ -19,7 +19,6 @@ class TrainerProtocol(Protocol): """ _batch_size: int - _epochs: int _patience: int _device: torch.device @@ -44,6 +43,18 @@ def train(self, *args: Any, **kwargs: Any) -> None: ... @property def epoch(self) -> int: ... + @property + def batch_size(self) -> int: ... + + @property + def train_n(self) -> int: ... + + @property + def val_n(self) -> int: ... + + @property + def test_n(self) -> int: ... + @property def device(self) -> torch.device: ... @@ -59,9 +70,19 @@ def best_model(self) -> Optional[torch.nn.Module]: ... def save_model( self, save_path: pathlib.Path, - file_name_prefix: Optional[str], - file_name_suffix: Optional[str], + file_name_prefix: str = 'generator', + file_name_suffix: Optional[str] = None, file_ext: str = '.pth', best_model: bool = True, ) -> Optional[List[pathlib.Path]]: ... + + def save_optimizer_state( + self, + save_path: pathlib.Path, + file_name_prefix: str = 'optimizer', + file_name_suffix: Optional[str] = None, + file_ext: str = '.pth', + recent: bool = True, + ) -> Optional[List[pathlib.Path]]: + ... diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py index f0a09fa..a706015 100644 --- a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py +++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py @@ -7,10 +7,12 @@ _get_latest_metric_value, ) from .save_model import save_model +from .save_optimizer import save_optimizer_state __all__ = [ "EarlyStopHelper", "_get_latest_metric_value", "save_model", + "save_optimizer_state", ] diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py new file mode 100644 index 0000000..ecdb510 --- /dev/null +++ b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py @@ -0,0 +1,36 @@ +from pathlib import Path +from typing import Optional, List + +import torch + +from ..trainer_protocol import TrainerProtocol + + +def save_optimizer_state( + trainer: 'TrainerProtocol', + save_path: Path, + file_name_prefix: str = 'optimizer', + file_name_suffix: Optional[str] = None, + file_ext: str = '.pth', +) -> List[Path]: + + if file_name_suffix is None: + file_name_suffix = f"{trainer.epoch}" + + optimizer = trainer.optimizer + + if optimizer is None: + return [] + + save_file = save_path / f"{file_name_prefix}_{file_name_suffix}{file_ext}" + + torch.save( + optimizer.state_dict(), + save_file + ) + + if save_file.exists(): + return [save_file] + + return [] + diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py index 8933283..b58d023 100644 --- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py +++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py @@ -15,7 +15,9 @@ AutoLossGroupConfigLogger, AutoModelConfigLogger, AutoOptimizerConfigLogger, + AutoTrainerLogger ) +from .logger_utils import _log_artifact, _save_and_log_trainer_artifacts from .callbacks.LoggerCallback import ( AbstractLoggerCallback, log_type @@ -145,6 +147,7 @@ def __init__( self._model_config_logger = AutoModelConfigLogger(self) self._optimizer_config_logger = AutoOptimizerConfigLogger(self) self._loss_group_config_logger = AutoLossGroupConfigLogger(self) + self._trainer_logger = AutoTrainerLogger(self) return None @@ -208,7 +211,7 @@ def on_train_start(self): self._model_config_logger.log_model_configs(self.trainer) self._optimizer_config_logger.log_optimizer_configs(self.trainer) self._loss_group_config_logger.log_loss_group_configs(self.trainer) - + self._trainer_logger.log_trainer_config(self.trainer) for callback in self.callbacks: # TODO consider if we want hasattr checks @@ -248,16 +251,10 @@ def on_epoch_end(self): if self._save_model_every_n_epochs is not None: if self.trainer.epoch % self._save_model_every_n_epochs == 0: - self._save_model_weights( - artifact_path='weights', - best_model=False - ) + _save_and_log_trainer_artifacts(self.trainer, best_model=False) if self._save_best_model: - self._save_model_weights( - artifact_path='weights', - best_model=True - ) + _save_and_log_trainer_artifacts(self.trainer, best_model=True) # Call on_epoch_end for all registered callbacks for callback in self.callbacks: @@ -279,16 +276,10 @@ def on_train_end(self): """ # Save weights to a temporary directory and log artifacts if self._save_model_at_train_end: - self._save_model_weights( - artifact_path='weights', - best_model=False - ) + _save_and_log_trainer_artifacts(self.trainer, best_model=False) if self._save_best_model: - self._save_model_weights( - artifact_path='weights', - best_model=True - ) + _save_and_log_trainer_artifacts(self.trainer, best_model=True) for callback in self.callbacks: if hasattr(callback, 'on_train_end'): @@ -323,14 +314,14 @@ def end_run(self): print("No active MLflow run to end.") """ - Exposed? logging methods + Exposed logging methods """ def log_artifact( - self, - tag: str, - file_path: pathlib.Path, - stage: Optional[str] = None - ): + self, + tag: str, + file_path: pathlib.Path, + stage: Optional[str] = None + ): """ Log an artifact to MLflow. @@ -339,29 +330,11 @@ def log_artifact( :param stage: Optional stage to categorize the artifact, defaults to None. :raises TypeError: If file_path is not a pathlib.Path instance. """ - - if not isinstance(file_path, pathlib.Path): - raise TypeError("file_path must be a pathlib.Path instance.") - - artifact_path = '' - artifact_ext = file_path.suffix.lower() - if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']: - # log as plot artifact - artifact_path += 'plots/' - elif artifact_ext in ['.pth', '.pt']: - # log as model artifact - artifact_path += 'weights/' - else: - # log as generic artifact - artifact_path += 'artifacts/' - - if stage is not None: - artifact_path += f"{stage}/" - artifact_path += f"{tag}" - - mlflow.log_artifact( - str(file_path), - artifact_path=artifact_path + log_subdirs = [stage] if stage is not None else [] + log_subdirs = log_subdirs + [tag] if tag is not None else log_subdirs + _log_artifact( + file_path=file_path, + artifact_subdirs=log_subdirs ) def log_metric( @@ -464,31 +437,6 @@ def _log_callback_output( continue # raise TypeError("Unsupported callback return type for logging.") - def _save_model_weights( - self, - prefix: Optional[str] = None, - suffix: Optional[str] = None, - artifact_path: str = "weights", - best_model: bool = True - ): - with tempfile.TemporaryDirectory() as tmpdirname: - - tmpdirpath = pathlib.Path(tmpdirname) - - saved_file_paths = self.trainer.save_model( - save_path=tmpdirpath, - file_name_prefix=prefix, - file_name_suffix=suffix, - file_ext='.pth', - best_model=best_model - ) - - for saved_file_path in (saved_file_paths or []): - mlflow.log_artifact( - str(saved_file_path), - artifact_path=artifact_path - ) - def log_config( self, tag: str, @@ -550,7 +498,6 @@ def __check_trainer_bound( if self.trainer is None: raise RuntimeError("No trainer bound to logger. Cannot access trainer attributes.") - def get_epoch( self ) -> int: @@ -595,48 +542,8 @@ def get_model( @property def run_id(self): return self._run_id - - """ - Unimplemented helper that might be useful - """ - def _log_dict_as_yaml( - self, - dict: Dict[str, Any], - ): - """ - Log a dictionary as a YAML file in MLflow. - - :param dict: The dictionary to log. - """ - - # TODO implement this method to convert dict to YAML and log it - raise NotImplementedError( - "log_dict_as_yaml method is not implemented. " - "Please implement this method to log dictionary as YAML." - ) - - def _log_dict_as_param( - self - ): - """ - Log a dictionary as parameters in MLflow. - - :return: None - :raises NotImplementedError: If the method is not implemented. - - This method is intended to log a dictionary as parameters in MLflow. - It is currently not implemented and raises a NotImplementedError. - """ - - # 1. DO flatten dict - - # 2. DO mlflow log - - raise NotImplementedError( - "log_dict_as_param method is not implemented. " - "Please implement this method to log dictionary as parameters." - ) + """ Overridden destructor method to ensure MLflow run is ended """ diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py index 2860937..2625d41 100644 --- a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py +++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py @@ -1,9 +1,11 @@ from .loss_group_config_logger import AutoLossGroupConfigLogger from .model_config_logger import AutoModelConfigLogger from .optimizer_config_logger import AutoOptimizerConfigLogger +from .trainer_config_logger import AutoTrainerLogger __all__ = [ "AutoModelConfigLogger", "AutoOptimizerConfigLogger", "AutoLossGroupConfigLogger", + "AutoTrainerLogger", ] diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py index fa3f344..c95537e 100644 --- a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py +++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py @@ -1,4 +1,5 @@ from typing import Any, Dict, List, Optional +import inspect import mlflow from torch.optim import Optimizer @@ -44,13 +45,22 @@ def log_optimizer_configs( continue try: + + defaults = dict(optimizer.defaults) + init_signature = inspect.signature(optimizer.__class__.__init__) + valid_params = init_signature.parameters.keys() + opt_config: Optional[Dict[str, Any]] = { "class_path": ( f"{optimizer.__class__.__module__}." f"{optimizer.__class__.__name__}" ), - "defaults": dict(optimizer.defaults), + "defaults": defaults, + "init": { + k: v for k, v in defaults.items() if k in valid_params + } } + except Exception as e: print(f"Could not get optimizer config for logging: {e}") opt_config = None diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py new file mode 100644 index 0000000..cf6c292 --- /dev/null +++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py @@ -0,0 +1,35 @@ +from typing import Any, Optional + +from ...trainers.trainer_protocol import TrainerProtocol + + +class AutoTrainerLogger: + """ + Auto-log trainer metadata to MLflow. + """ + + def __init__(self, logger: Any) -> None: + self._logger = logger + + def log_trainer_config(self, trainer: Optional[TrainerProtocol]) -> None: + + if trainer is None: + return + + config = { + "class_path": f"{trainer.__class__.__module__}.{trainer.__class__.__name__}", + "device": str(trainer.device), # device used for training + "batch_size": trainer.batch_size, # batch size used for training + "train_n": trainer.train_n, + "val_n": trainer.val_n, + "test_n": trainer.test_n, + } + + try: + self._logger.log_config( + tag="trainer", + config=config, + stage=None, + ) + except Exception as e: + print(f"Could not log trainer config: {e}") diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py new file mode 100644 index 0000000..e7edb5e --- /dev/null +++ b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py @@ -0,0 +1,6 @@ +from .log_artifacts import _log_artifact, _save_and_log_trainer_artifacts + +__all__ = [ + "_log_artifact", + "_save_and_log_trainer_artifacts", +] diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py new file mode 100644 index 0000000..2bf371a --- /dev/null +++ b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py @@ -0,0 +1,84 @@ +from typing import Optional, List +import pathlib +import tempfile + +import mlflow + +from virtual_stain_flow.trainers.trainer_protocol import TrainerProtocol + + +def _log_artifact( + file_path: pathlib.Path, + artifact_path: Optional[str] = None, + artifact_subdirs: Optional[List[str]] = None +) -> None: + """ + Logs a single artifact to MLflow. + + :param file_path: The path to the file to log as an artifact. + :param artifact_path: Optional artifact path within the MLflow run, defaults to None. + :param artifact_subdirs: Optional list of subdirectories to include in the artifact path, defaults to None. + :raises TypeError: If file_path is not a pathlib.Path instance. + """ + + if not isinstance(file_path, pathlib.Path): + raise TypeError("file_path must be a pathlib.Path instance.") + + if artifact_path is None: + artifact_ext = file_path.suffix.lower() + if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']: + # log as plot artifact + artifact_path = 'plots' + elif artifact_ext in ['.pth', '.pt']: + # log as model artifact + artifact_path = 'weights' + else: + # log as generic artifact + artifact_path = 'artifacts' + + artifact_subdirs = [] if artifact_subdirs is None else artifact_subdirs + + path_parts = [artifact_path, *artifact_subdirs] + clean_parts = [ + part.replace('\\', '/').strip('/') + for part in path_parts + if part and part.strip('/\\') + ] + + # Build a lexical POSIX path for MLflow without resolving it on the host OS. + artifact_path = str(pathlib.PurePosixPath(*clean_parts)) if clean_parts else '' + + mlflow.log_artifact(str(file_path), artifact_path=artifact_path) + + +def _save_and_log_trainer_artifacts( + trainer: "TrainerProtocol", + best_model: bool = True, +) -> None: + """ + Saves the trainer's model and optimizer state to temporary files and logs + those files as MLflow artifacts. + The most recent optimizer state is saved and logged regardless of the best_model flag. + + :param trainer: The trainer instance adhering to TrainerProtocol. + :param best_model: Whether to save and log the best model, defaults to True. + :raises TypeError: If the provided trainer does not adhere to TrainerProtocol. + """ + + if not isinstance(trainer, TrainerProtocol): + raise TypeError("The provided trainer must adhere to the TrainerProtocol.") + + with tempfile.TemporaryDirectory() as tmpdirname: + + tmpdirpath = pathlib.Path(tmpdirname) + saved_model_paths = trainer.save_model( + save_path=tmpdirpath, best_model=best_model + ) + for saved_model_path in (saved_model_paths or []): + _log_artifact(saved_model_path, artifact_path='weights') + + saved_optimizer_paths = trainer.save_optimizer_state( + save_path=tmpdirpath, recent=True + ) + for saved_optimizer_path in (saved_optimizer_paths or []): + _log_artifact(saved_optimizer_path, artifact_path='optimizer') diff --git a/tests/conftest.py b/tests/conftest.py index 5842619..9c42267 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -388,7 +388,7 @@ def evaluate_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict: 'loss_b': torch.tensor(0.2), } - def save_model(self, save_path, file_name_prefix=None, file_name_suffix=None, + def save_model(self, save_path, file_name_prefix='generator', file_name_suffix=None, file_ext='.pth', best_model=True): return None diff --git a/tests/datasets/ds_engine/test_manifest.py b/tests/datasets/ds_engine/test_manifest.py index 168071f..d174a78 100644 --- a/tests/datasets/ds_engine/test_manifest.py +++ b/tests/datasets/ds_engine/test_manifest.py @@ -362,6 +362,27 @@ def test_from_config_missing_manifest(self): class TestDatasetManifestSerialization: """Test suite for DatasetManifest serialization methods.""" + def test_config_round_trip(self): + """Test serialization converts paths to strings and restores Paths.""" + manifest = DatasetManifest( + file_index=pd.DataFrame({ + "channel1": [Path("/path/to/img1.tif")], + "channel2": ["/path/to/img2.tif"], + }) + ) + + config = manifest.to_config() + restored = DatasetManifest.from_config(config) + + assert config["file_index"] == [{ + "channel1": "/path/to/img1.tif", + "channel2": "/path/to/img2.tif", + }] + assert restored.file_index.to_dict(orient="records") == [{ + "channel1": Path("/path/to/img1.tif"), + "channel2": Path("/path/to/img2.tif"), + }] + def test_from_config_missing_file_index(self): """Test DatasetManifest.from_config raises ValueError when file_index is missing.""" config = {"pil_image_mode": "I;16", "file_index": None} diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py index 2cb9099..4474603 100644 --- a/tests/trainers/test_abstract_trainer.py +++ b/tests/trainers/test_abstract_trainer.py @@ -338,6 +338,37 @@ def test_train_epoch_with_large_batch_count(self, minimal_model, minimal_optimiz class TestDataSplitting: """Test that AbstractTrainer correctly handles dataset splitting.""" + + def test_init_rejects_none_epoch( + self, minimal_model, minimal_optimizer, train_dataloader + ): + with pytest.raises(TypeError, match="epoch must be an integer"): + MinimalTrainerRealization( + model=minimal_model, + optimizer=minimal_optimizer, + train_loader=train_dataloader, + epoch=None, + device=torch.device('cpu') + ) + + def test_init_with_unsized_dataset_records_unknown_size( + self, minimal_model, minimal_optimizer + ): + from torch.utils.data import DataLoader, IterableDataset + + class UnsizedDataset(IterableDataset): + def __iter__(self): + yield torch.randn(4), torch.randn(2) + + train_loader = DataLoader(UnsizedDataset(), batch_size=1) + trainer = MinimalTrainerRealization( + model=minimal_model, + optimizer=minimal_optimizer, + train_loader=train_loader, + device=torch.device('cpu') + ) + + assert trainer.train_n is None def test_init_with_dataset_creates_loaders(self, minimal_model, minimal_optimizer, dataset_for_splitting): """Verify that providing a dataset creates train/val/test loaders.""" @@ -697,8 +728,9 @@ def test_batch_size_property_with_loader_init(self, minimal_model, minimal_optim device=torch.device('cpu') ) - # When providing loaders, batch_size is set to None - assert trainer.batch_size is None + # When providing loaders, batch_size is inferred from the train loader + # should therefore match + assert trainer.batch_size is train_dataloader.batch_size def test_batch_size_property_default_value(self, minimal_model, minimal_optimizer, dataset_for_splitting): """Verify that batch_size property uses default value when not specified.""" diff --git a/tests/trainers/test_abstract_trainer_save_artifacts.py b/tests/trainers/test_abstract_trainer_save_artifacts.py new file mode 100644 index 0000000..8f5ec49 --- /dev/null +++ b/tests/trainers/test_abstract_trainer_save_artifacts.py @@ -0,0 +1,82 @@ +"""Tests for AbstractTrainer save-related class methods.""" + +import torch + +from virtual_stain_flow.trainers.AbstractTrainer import AbstractTrainer + + +class MinimalArtifactTrainer(AbstractTrainer): + """Concrete trainer realization to exercise base save methods.""" + + def train_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict: + return {"loss": torch.tensor(0.0)} + + def evaluate_step(self, inputs: torch.Tensor, targets: torch.Tensor) -> dict: + return {"loss": torch.tensor(0.0)} + + +def test_save_model_writes_current_model_file( + mock_model_with_save, + mock_optimizer, + train_dataloader, + val_dataloader, + tmp_path, +): + trainer = MinimalArtifactTrainer( + model=mock_model_with_save, + optimizer=mock_optimizer, + train_loader=train_dataloader, + val_loader=val_dataloader, + device=torch.device("cpu"), + ) + + saved_paths = trainer.save_model(save_path=tmp_path, best_model=False) + + assert isinstance(saved_paths, list) + assert len(saved_paths) == 1 + assert saved_paths[0].exists() + assert saved_paths[0].name == "generator_weights_0.pth" + + +def test_save_optimizer_state_writes_recent_file( + mock_model_with_save, + mock_optimizer, + train_dataloader, + val_dataloader, + tmp_path, +): + trainer = MinimalArtifactTrainer( + model=mock_model_with_save, + optimizer=mock_optimizer, + train_loader=train_dataloader, + val_loader=val_dataloader, + device=torch.device("cpu"), + ) + + saved_paths = trainer.save_optimizer_state(save_path=tmp_path, recent=True) + + assert isinstance(saved_paths, list) + assert len(saved_paths) == 1 + assert saved_paths[0].exists() + assert saved_paths[0].name == "optimizer_recent.pth" + + +def test_save_optimizer_state_non_recent_not_supported( + mock_model_with_save, + mock_optimizer, + train_dataloader, + val_dataloader, + tmp_path, +): + trainer = MinimalArtifactTrainer( + model=mock_model_with_save, + optimizer=mock_optimizer, + train_loader=train_dataloader, + val_loader=val_dataloader, + device=torch.device("cpu"), + ) + + import pytest + + with pytest.raises(NotImplementedError, match="non-recent optimizer states"): + trainer.save_optimizer_state(save_path=tmp_path, recent=False) diff --git a/tests/trainers/test_logging_trainer.py b/tests/trainers/test_logging_trainer.py index c4c455b..0250272 100644 --- a/tests/trainers/test_logging_trainer.py +++ b/tests/trainers/test_logging_trainer.py @@ -559,6 +559,14 @@ def test_train_logged_metrics_have_correct_steps(self, conv_trainer, dummy_logge epochs = 3 conv_trainer.train(logger=dummy_logger, epochs=epochs, verbose=False) - # Check that step numbers are within expected range (0-indexed) - for metric in dummy_logger.logged_metrics: - assert 0 <= metric['step'] < epochs + assert {metric['step'] for metric in dummy_logger.logged_metrics} == {1, 2, 3} + + def test_resumed_training_logs_cumulative_epoch_steps( + self, conv_trainer, dummy_logger + ): + """Metric steps stay aligned with callbacks after resumed training.""" + conv_trainer.epoch = 5 + + conv_trainer.train(logger=dummy_logger, epochs=2, verbose=False) + + assert {metric['step'] for metric in dummy_logger.logged_metrics} == {6, 7} diff --git a/tests/trainers/trainer_utils/test_save_model.py b/tests/trainers/trainer_utils/test_save_model.py new file mode 100644 index 0000000..be65008 --- /dev/null +++ b/tests/trainers/trainer_utils/test_save_model.py @@ -0,0 +1,23 @@ +"""Standalone tests for save_model helper.""" + +from types import SimpleNamespace + +from virtual_stain_flow.trainers.trainer_utils.save_model import save_model + + +def test_save_model_returns_empty_when_target_model_is_none(tmp_path): + trainer = SimpleNamespace(model=None, best_model=None, epoch=2) + + saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=True) + + assert saved_paths == [] + + +def test_save_model_saves_current_model_with_default_name(mock_model_with_save, tmp_path): + trainer = SimpleNamespace(model=mock_model_with_save, best_model=None, epoch=3) + + saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=False) + + assert len(saved_paths) == 1 + assert saved_paths[0].exists() + assert saved_paths[0].name == "generator_weights_3.pth" diff --git a/tests/trainers/trainer_utils/test_save_optimizer.py b/tests/trainers/trainer_utils/test_save_optimizer.py new file mode 100644 index 0000000..dac7d52 --- /dev/null +++ b/tests/trainers/trainer_utils/test_save_optimizer.py @@ -0,0 +1,23 @@ +"""Standalone tests for save_optimizer_state helper.""" + +from types import SimpleNamespace + +from virtual_stain_flow.trainers.trainer_utils.save_optimizer import save_optimizer_state + + +def test_save_optimizer_state_returns_empty_for_missing_optimizer(tmp_path): + trainer = SimpleNamespace(optimizer=None, epoch=7) + + saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path) + + assert saved_paths == [] + + +def test_save_optimizer_state_saves_with_default_name(minimal_optimizer, tmp_path): + trainer = SimpleNamespace(optimizer=minimal_optimizer, epoch=7) + + saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path) + + assert len(saved_paths) == 1 + assert saved_paths[0].exists() + assert saved_paths[0].name == "optimizer_7.pth"