From ad3cbf4ee211a08644cb47e258b688b3520eeb29 Mon Sep 17 00:00:00 2001 From: Bowen Li Date: Wed, 9 Sep 2026 16:25:50 +0100 Subject: [PATCH 1/2] Implement bootstrap ensemble strategy and add MLP bootstrap example configs --- configs/mlp_bootstrap.yaml | 81 ++++++ configs/mlp_bootstrap_infer.yaml | 23 ++ .../schemas/strategies/bootstrap.schema.yaml | 18 +- xanesnet/serialization/runtime_contracts.py | 12 +- xanesnet/strategies/bootstrap.py | 257 +++++++++++++++--- 5 files changed, 339 insertions(+), 52 deletions(-) create mode 100644 configs/mlp_bootstrap.yaml create mode 100644 configs/mlp_bootstrap_infer.yaml diff --git a/configs/mlp_bootstrap.yaml b/configs/mlp_bootstrap.yaml new file mode 100644 index 00000000..5312c3e3 --- /dev/null +++ b/configs/mlp_bootstrap.yaml @@ -0,0 +1,81 @@ +# Example configuration for MLP bootstrap ensemble training with toy data. + +seed: 2025 +device: cpu + +datasource: + datasource_type: pmgjson + json_path: ./data/toy_data/ + spectrum_key: "XANES" + +dataset: + # general: + dataset_type: descriptor + root: ./data/processed/toy_data_mlp_bootstrap/ # Where should the processed data be stored + preload: True # Preload dataset into RAM if True; otherwise load on-the-fly + skip_prepare: False + split_ratios: [0.8, 0.2] + # params: + # descriptors: + descriptors: + - descriptor_type: wacsf + r_min: 1.0 + r_max: 6.0 + n_g2: 16 + n_g4: 32 + +encodings: + - encoding_type: identity + +model: + # general: + model_type: mlp + # params: + in_size: auto + out_size: auto + hidden_size: 256 + dropout: 0.1 + num_hidden_layers: 3 + shrink_rate: 0.5 + activation: prelu + +trainer: + # general: + trainer_type: basic + batch_size: 4 + shuffle: True + drop_last: False + num_workers: 0 + # params: + epochs: 20 + learning_rate: 0.001 + optimizer: Adam + max_norm: null + validation_interval: 10 + lr_warmup: True + warmup_steps: 500 + loss: + - loss_type: mse + regularizer: + regularizer_type: none + lr_scheduler: + lr_scheduler_type: linear + start_factor: 1.0 + end_factor: 0.1 + total_iters: 40 + early_stopper: + early_stopper_type: basic + patience: 25 + min_delta: 0.001 + restore_best: True + +strategy: + # general: + strategy_type: bootstrap + weight_init: xavier_uniform + bias_init: zeros + checkpoint_interval: 25 + # params: + n_models: 5 + sample_fraction: 1.0 + seeds: [101, 202, 303, 404, 505] diff --git a/configs/mlp_bootstrap_infer.yaml b/configs/mlp_bootstrap_infer.yaml new file mode 100644 index 00000000..3db57e16 --- /dev/null +++ b/configs/mlp_bootstrap_infer.yaml @@ -0,0 +1,23 @@ +# Example configuration for MLP bootstrap ensemble inference with toy data. + +seed: 2025 +device: cuda + +datasource: + datasource_type: pmgjson + json_path: ./data/toy_data/ + spectrum_key: "XANES" + +dataset: + root: ./data/processed/toy_data_mlp_bootstrap/ # Where should the processed data be stored + preload: True # Preload dataset into RAM if True; otherwise load on-the-fly + skip_prepare: False + +inferencer: + inferencer_type: ensemble + batch_size: 4 + shuffle: False + drop_last: False + num_workers: 0 + buffer_size: 1000 + model_device_policy: sequential diff --git a/xanesnet/schemas/strategies/bootstrap.schema.yaml b/xanesnet/schemas/strategies/bootstrap.schema.yaml index 82269087..da53bc93 100644 --- a/xanesnet/schemas/strategies/bootstrap.schema.yaml +++ b/xanesnet/schemas/strategies/bootstrap.schema.yaml @@ -1,7 +1,6 @@ $schema: https://json-schema.org/draft/2020-12/schema $id: bootstrap.schema.yaml title: 'XANESNET strategy schema: bootstrap' -$comment: Registered in the strategy schemas, but the implementation currently raises NotImplementedError. $defs: strategyBootstrap: type: object @@ -33,5 +32,20 @@ $defs: strategy_type: const: bootstrap description: Strategy identifier. - description: Bootstrap ensemble strategy. + n_models: + type: integer + minimum: 1 + default: 5 + description: Number of bootstrap ensemble members. + sample_fraction: + type: number + exclusiveMinimum: 0 + default: 1.0 + description: Fraction of the training subset resampled for each bootstrap member. + seeds: + type: array + items: + type: integer + description: Optional per-model random seeds for resampling and weight initialization. + description: Sequential bootstrap-ensemble training and aggregate inference strategy. description: Strategy schemas backed by ``StrategyRegistry`` classes. diff --git a/xanesnet/serialization/runtime_contracts.py b/xanesnet/serialization/runtime_contracts.py index a6f740ac..8cb1f2fb 100644 --- a/xanesnet/serialization/runtime_contracts.py +++ b/xanesnet/serialization/runtime_contracts.py @@ -110,20 +110,20 @@ def _require_concrete_inference_encodings(config: ConfigRaw) -> None: ) -def _require_ensemble_inferencer_for_deep_ensemble(config: ConfigRaw) -> None: - """Validate deep-ensemble inference runner selection. +def _require_ensemble_inferencer_for_ensemble_strategies(config: ConfigRaw) -> None: + """Validate ensemble-strategy inference runner selection. Args: config: Schema-valid merged inference configuration. Raises: - ConfigError: If a deep-ensemble strategy is paired with a non-ensemble + ConfigError: If an ensemble strategy is paired with a non-ensemble inferencer. """ strategy_type = _section_value(config, "strategy", "strategy_type") inferencer_type = _section_value(config, "inferencer", "inferencer_type") - if strategy_type == "deep_ensemble" and inferencer_type != "ensemble": - raise ConfigError("Inference strategy 'deep_ensemble' requires inferencer 'ensemble'.") + if strategy_type in {"deep_ensemble", "bootstrap"} and inferencer_type != "ensemble": + raise ConfigError(f"Inference strategy '{strategy_type}' requires inferencer 'ensemble'.") def _section_value(config: ConfigRaw, section: str, key: str) -> Any: @@ -169,7 +169,7 @@ def _auto_token_paths(value: Any, path: tuple[str, ...]) -> Iterator[str]: _require_registered_batch_processor, _require_concrete_inference_model, _require_concrete_inference_encodings, - _require_ensemble_inferencer_for_deep_ensemble, + _require_ensemble_inferencer_for_ensemble_strategies, ), "analyze": (), } diff --git a/xanesnet/strategies/bootstrap.py b/xanesnet/strategies/bootstrap.py index 6cc017da..a7640d40 100644 --- a/xanesnet/strategies/bootstrap.py +++ b/xanesnet/strategies/bootstrap.py @@ -18,16 +18,25 @@ # Citations: # ... -"""Bootstrap ensemble strategy for XANESNET (placeholder).""" +"""Bootstrap ensemble training and inference strategy for XANESNET.""" +import copy +import logging +import random from pathlib import Path +from typing import Any +import numpy as np import torch +from torch.utils.data import Subset from xanesnet.datasets import Dataset from xanesnet.encodings import SpectraEncoding -from xanesnet.models import Model +from xanesnet.models import Model, ModelRegistry +from xanesnet.runners.inferencers import InferencerRegistry +from xanesnet.runners.trainers import TrainerRegistry from xanesnet.serialization.config import Config +from xanesnet.serialization.tensorboard import tb_logger from .base import Strategy from .registry import StrategyRegistry @@ -35,21 +44,29 @@ @StrategyRegistry.register("bootstrap") class Bootstrap(Strategy): - """Bootstrap ensemble strategy. + """Sequential bootstrap-ensemble training and aggregate inference strategy. - Note: - Not yet implemented. - Will be implemented later (low priority). - TODO: implement bootstrap sampling for model ensembling. + The strategy owns ``n_models`` independent model instances with identical + architecture. During training, each member is trained on a bootstrap + resample of the training subset (sampling with replacement). During + inference, all member models are evaluated on the same batches and their + predictions are reduced to mean and energy/channel-wise standard deviation + by the ensemble inferencer. Args: - strategy_type: Strategy identifier. + strategy_type: Registry key identifying this strategy type. dataset: Dataset used for training or inference. model_config: Configuration for the model. - encoding: Composed spectra encoding. + encoding: Composed spectra encoding forwarded to the trainers and + inferencer. weight_init: Weight initialization scheme name. weight_init_params: Additional weight-initializer parameters. bias_init: Bias initialization scheme name. + n_models: Number of bootstrap ensemble members. + sample_fraction: Fraction of the training subset drawn for each + bootstrap resample. + seeds: Optional per-model random seeds used for resampling and weight + initialization. When omitted, seeds are sampled automatically. checkpoint_dir: Directory for checkpoints, or ``None``. checkpoint_interval: Epoch interval between checkpoints, or ``None``. tensorboard_dir: Directory for TensorBoard event files, or ``None``. @@ -66,13 +83,16 @@ def __init__( weight_init: str, weight_init_params: Config, bias_init: str, + n_models: int, checkpoint_dir: str | Path | None, checkpoint_interval: int | None, tensorboard_dir: str | Path | None, + sample_fraction: float = 1.0, + seeds: list[int] | None = None, trainer_config: Config | None = None, inferencer_config: Config | None = None, ) -> None: - """Initialize the placeholder bootstrap strategy.""" + """Initialize the bootstrap ensemble strategy.""" super().__init__( strategy_type, dataset, @@ -88,100 +108,249 @@ def __init__( inferencer_config, ) - def setup_models(self) -> None: - """Raise because bootstrap model setup is not implemented. + self.n_models = n_models + self.sample_fraction = sample_fraction + if seeds is None: + self.seeds = random.sample(range(1000), n_models) + elif len(seeds) != n_models: + raise ValueError(f"Expected {n_models} bootstrap seeds, got {len(seeds)}.") + else: + self.seeds = seeds - Raises: - NotImplementedError: Always. + self.models: list[Model] = [] + self.trainers: list[Any | None] = [] + self.inferencer: Any | None = None + + def _bootstrap_dataset(self, model_idx: int) -> Dataset: + """Return a dataset copy with a bootstrap-resampled training subset. + + Args: + model_idx: Bootstrap member index used to select the resampling seed. + + Returns: + A shallow copy of ``self.dataset`` whose training subset contains + a bootstrap resample of the original training indices. """ - raise NotImplementedError("Not implemented!") # TODO Implement + train_indices = self.dataset.get_subset_indices(0) + if train_indices is None: + train_indices = list(range(len(self.dataset))) + + rng = np.random.default_rng(self.seeds[model_idx]) + n_samples = len(train_indices) + sample_size = int(n_samples * self.sample_fraction) + bootstrap_positions = rng.choice(n_samples, size=sample_size, replace=True) + bootstrap_indices = [train_indices[i] for i in bootstrap_positions] + + dataset_boot = copy.copy(self.dataset) + subsets: list[Subset] = [Subset(self.dataset, bootstrap_indices)] + valid_subset = self.dataset.valid_subset + if valid_subset is not None: + subsets.append(valid_subset) + dataset_boot._subsets = subsets + + return dataset_boot + + def setup_models(self) -> None: + """Instantiate ``n_models`` independent model copies from ``model_config``.""" + model_type = self.model_config.get_str("model_type") + model_cls = ModelRegistry.get(model_type) + + self.models = [] + for model_idx in range(self.n_models): + logging.info(f"Initializing bootstrap model {model_idx + 1}/{self.n_models}: {model_type}") + self.models.append(model_cls(**self.model_config.as_kwargs())) def init_model_weights(self) -> None: - """Raise because bootstrap weight initialization is not implemented. + """Apply configured weight and bias initialization to every model. + + Each member is initialized with its own bootstrap seed so that weight + draws are reproducible and distinct across members. Raises: - NotImplementedError: Always. + ValueError: If ``setup_models`` has not been called. """ - raise NotImplementedError("Not implemented!") # TODO Implement + if len(self.models) == 0: + raise ValueError("Cannot initialize model weights because models are not initialized.") + + logging.info(f"Initializing weights with '{self.weight_init}' and bias with '{self.bias_init}'") + for model_idx, model in enumerate(self.models): + logging.info(f"Initializing bootstrap model {model_idx + 1}/{self.n_models} weights.") + torch.manual_seed(self.seeds[model_idx]) + model.init_weights(self.weight_init, self.bias_init, **self.weight_init_params.as_kwargs()) def set_state_dicts(self, state_dicts: list[dict]) -> None: - """Raise because bootstrap state-dict loading is not implemented. + """Load one state dictionary into each bootstrap member. Args: - state_dicts: State dictionaries that would be loaded into managed models. + state_dicts: State dictionaries to load, one per model. Raises: - NotImplementedError: Always. + ValueError: If models are not initialized or the number of state + dictionaries does not match the number of models. """ - raise NotImplementedError("Not implemented!") # TODO Implement + if len(self.models) == 0: + raise ValueError("Cannot load state dicts because models are not initialized.") + if len(state_dicts) != len(self.models): + raise ValueError(f"Expected {len(self.models)} state dicts, got {len(state_dicts)}.") + + for model, state_dict in zip(self.models, state_dicts, strict=True): + model.load_state_dict(state_dict) def setup_trainers(self, device: str | torch.device) -> None: - """Raise because bootstrap trainer setup is not implemented. + """Instantiate one trainer per bootstrap member. + + Each trainer receives a dataset copy whose training subset is a + bootstrap resample of the original training data. + + Must be called after ``setup_models`` and ``setup_checkpointer``. Args: - device: Target device for training. + device: The device on which training will be performed. Raises: - NotImplementedError: Always. + ValueError: If models, trainer config, or checkpointer are not initialized. """ - raise NotImplementedError("Not implemented!") # TODO Implement + if len(self.models) == 0: + raise ValueError("Cannot setup trainers because models are not initialized.") + if self.trainer_config is None: + raise ValueError("Can not setup trainers because there is no trainer config.") + if self.checkpointer is None: + raise ValueError("Can not setup trainers because checkpointer is not instantiated.") + + trainer_type = self.trainer_config.get_str("trainer_type") + trainer_cls = TrainerRegistry.get(trainer_type) + + self.trainers = [] + for model_idx, model in enumerate(self.models): + logging.info(f"Initializing trainer {model_idx + 1}/{self.n_models}: {trainer_type}") + dataset_boot = self._bootstrap_dataset(model_idx) + trainer = trainer_cls( + **self.trainer_config.as_kwargs(), + dataset=dataset_boot, + model=model, + device=device, + checkpointer=self.checkpointer, + encoding=self.encoding, + ) + self.trainers.append(trainer) def run_training(self) -> list[Model]: - """Raise because bootstrap training is not implemented. + """Train all bootstrap members sequentially and return them. + + Must be called after ``setup_trainers``. Returns: - Never returns normally. + List of trained bootstrap member models. Trainer instances are + released after their corresponding member finishes to avoid keeping + optimizer state alive during later member training. Raises: - NotImplementedError: Always. + ValueError: If trainers or models are not initialized. """ + if len(self.models) == 0: + raise ValueError("Cannot run training because models are not initialized.") + if len(self.trainers) != len(self.models): + raise ValueError("Cannot run training because trainers are not initialized for every model.") + super().run_training() - raise NotImplementedError("Not implemented!") # TODO Implement + assert self.checkpointer is not None + + for model_idx, trainer in enumerate(self.trainers): + if trainer is None: + raise ValueError("Cannot run training because trainers are not initialized for every model.") + + logging.info(f"Training bootstrap model {model_idx + 1}/{self.n_models}.") + self.checkpointer.new_model() + + try: + if self.tensorboard_dir is not None: + tb_logger.new_run(Path(self.tensorboard_dir) / f"model_{model_idx}") + + trainer.train() + finally: + tb_logger.close() + self.models[model_idx].to(torch.device("cpu")) + self.trainers[model_idx] = None + + return self.models def setup_inferencers(self, device: str | torch.device) -> None: - """Raise because bootstrap inferencer setup is not implemented. + """Instantiate the ensemble inferencer for all loaded models. + + Must be called after ``setup_models`` and ``set_state_dicts``. Args: - device: Target device for inference. + device: The device on which inference will be performed. Raises: - NotImplementedError: Always. + ValueError: If models or inferencer config are not initialized. """ - raise NotImplementedError("Not implemented!") # TODO Implement + if len(self.models) == 0: + raise ValueError("Can not setup inferencers because models are not initialized.") + if self.inferencer_config is None: + raise ValueError("Can not setup inferencers because there is no inferencer config.") + + logging.info("Initializing inferencer: ensemble") + + inferencer_kwargs = self.inferencer_config.as_kwargs() + inferencer_kwargs["inferencer_type"] = "ensemble" + inferencer = InferencerRegistry.create( + "ensemble", + **inferencer_kwargs, + dataset=self.dataset, + models=self.models, + device=device, + encoding=self.encoding, + ) + + self.inferencer = inferencer def run_inference(self, predictions_save_path: str | Path | None) -> None: - """Raise because bootstrap inference is not implemented. + """Run aggregate ensemble inference and optionally save predictions. Args: - predictions_save_path: Destination directory for prediction output. + predictions_save_path: Directory in which to write prediction + output, or ``None`` to skip saving. Raises: - NotImplementedError: Always. + ValueError: If ``setup_inferencers`` has not been called. """ + if self.inferencer is None: + raise ValueError("Cannot run inference because the Inferencer is not initialized.") + super().run_inference(predictions_save_path) - raise NotImplementedError("Not implemented!") # TODO Implement + self.inferencer.infer(predictions_save_path) @property def model_signature(self) -> Config: - """Raise because bootstrap model signatures are not implemented. + """Return the shared model architecture signature. Returns: - Never returns normally. + A ``Config`` representing the bootstrap members' model signature. Raises: - NotImplementedError: Always. + ValueError: If ``setup_models`` has not been called. """ - raise NotImplementedError("Not implemented!") # TODO Implement + if len(self.models) == 0: + raise ValueError("Models are not initialized. Cannot retrieve signature.") + + return self.models[0].signature @property def signature(self) -> Config: - """Return the placeholder bootstrap strategy configuration. + """Return the strategy configuration as a ``Config``. Returns: A ``Config`` capturing the strategy configuration. """ signature = super().signature - signature.update_with_dict({}) + signature.update_with_dict( + { + "n_models": self.n_models, + "sample_fraction": self.sample_fraction, + "seeds": self.seeds, + } + ) return signature From bad85a35778cc86cf804da55df0cf8db0430d61d Mon Sep 17 00:00:00 2001 From: Bowen Li Date: Thu, 10 Sep 2026 09:51:35 +0100 Subject: [PATCH 2/2] Implement kfold strategy and add example config file --- configs/mlp_kfold.yaml | 80 ++++ xanesnet/schemas/strategies/kfold.schema.yaml | 49 +++ .../schemas/strategies/strategies.schema.yaml | 1 + .../strategies/strategy_types.schema.yaml | 1 + xanesnet/strategies/__init__.py | 2 + xanesnet/strategies/kfold.py | 377 ++++++++++++++++++ 6 files changed, 510 insertions(+) create mode 100644 configs/mlp_kfold.yaml create mode 100644 xanesnet/schemas/strategies/kfold.schema.yaml create mode 100644 xanesnet/strategies/kfold.py diff --git a/configs/mlp_kfold.yaml b/configs/mlp_kfold.yaml new file mode 100644 index 00000000..aaf6b6f9 --- /dev/null +++ b/configs/mlp_kfold.yaml @@ -0,0 +1,80 @@ +# Example configuration for MLP k-fold cross-validation training with toy data. + +seed: 2025 +device: cpu + +datasource: + datasource_type: pmgjson + json_path: ./data/toy_data/ + spectrum_key: "XANES" + +dataset: + # general: + dataset_type: descriptor + root: ./data/processed/toy_data_mlp_kfold/ # Where should the processed data be stored + preload: True # Preload dataset into RAM if True; otherwise load on-the-fly + skip_prepare: False + # params: + # descriptors: + descriptors: + - descriptor_type: wacsf + r_min: 1.0 + r_max: 6.0 + n_g2: 16 + n_g4: 32 + +encodings: + - encoding_type: identity + +model: + # general: + model_type: mlp + # params: + in_size: auto + out_size: auto + hidden_size: 256 + dropout: 0.1 + num_hidden_layers: 3 + shrink_rate: 0.5 + activation: prelu + +trainer: + # general: + trainer_type: basic + batch_size: 4 + shuffle: True + drop_last: False + num_workers: 0 + # params: + epochs: 20 + learning_rate: 0.001 + optimizer: Adam + max_norm: null + validation_interval: 10 + lr_warmup: True + warmup_steps: 500 + loss: + - loss_type: mse + regularizer: + regularizer_type: none + lr_scheduler: + lr_scheduler_type: linear + start_factor: 1.0 + end_factor: 0.1 + total_iters: 40 + early_stopper: + early_stopper_type: basic + patience: 25 + min_delta: 0.001 + restore_best: True + +strategy: + # general: + strategy_type: kfold + weight_init: xavier_uniform + bias_init: zeros + checkpoint_interval: 25 + # params: + n_splits: 3 + n_repeats: 1 + seed: 2025 diff --git a/xanesnet/schemas/strategies/kfold.schema.yaml b/xanesnet/schemas/strategies/kfold.schema.yaml new file mode 100644 index 00000000..f899de01 --- /dev/null +++ b/xanesnet/schemas/strategies/kfold.schema.yaml @@ -0,0 +1,49 @@ +$schema: https://json-schema.org/draft/2020-12/schema +$id: kfold.schema.yaml +title: 'XANESNET strategy schema: kfold' +$defs: + strategyKFold: + type: object + additionalProperties: false + required: + - strategy_type + properties: + weight_init: + $ref: ../components/component_types.schema.yaml#/$defs/weightInitName + default: default + description: Weight initialization scheme name. + weight_init_params: + type: object + default: {} + additionalProperties: true + $comment: Forwarded as kwargs to the selected torch.nn.init weight initializer. Examples include {a,b} for uniform_, + {mean,std} for normal_, {gain} for xavier_*, and {a,mode,nonlinearity} for kaiming_*. + description: Additional weight-initializer parameters. + bias_init: + $ref: ../components/component_types.schema.yaml#/$defs/biasInitName + default: zeros + description: Bias initialization scheme name. + checkpoint_interval: + type: + - integer + - 'null' + default: null + description: Epoch interval between checkpoints, or ``None``. + strategy_type: + const: kfold + description: Strategy identifier. + n_splits: + type: integer + minimum: 2 + default: 3 + description: Number of folds per repeat. + n_repeats: + type: integer + minimum: 1 + default: 1 + description: Number of repeated k-fold shuffles. + seed: + type: integer + description: Random seed used to shuffle samples before splitting. + description: Repeated k-fold cross-validation strategy returning the best fold model. +description: Strategy schemas backed by ``StrategyRegistry`` classes. diff --git a/xanesnet/schemas/strategies/strategies.schema.yaml b/xanesnet/schemas/strategies/strategies.schema.yaml index 7635b52e..89436a9b 100644 --- a/xanesnet/schemas/strategies/strategies.schema.yaml +++ b/xanesnet/schemas/strategies/strategies.schema.yaml @@ -7,6 +7,7 @@ $defs: - $ref: single.schema.yaml#/$defs/strategySingle - $ref: deep_ensemble.schema.yaml#/$defs/strategyDeepEnsemble - $ref: bootstrap.schema.yaml#/$defs/strategyBootstrap + - $ref: kfold.schema.yaml#/$defs/strategyKFold - $ref: snapshot_ensemble.schema.yaml#/$defs/strategySnapshotEnsemble description: Union of strategy configuration objects instantiated by ``StrategyRegistry``. description: Strategy schemas backed by ``StrategyRegistry`` classes. diff --git a/xanesnet/schemas/strategies/strategy_types.schema.yaml b/xanesnet/schemas/strategies/strategy_types.schema.yaml index d76c390d..b40af9d8 100644 --- a/xanesnet/schemas/strategies/strategy_types.schema.yaml +++ b/xanesnet/schemas/strategies/strategy_types.schema.yaml @@ -8,6 +8,7 @@ $defs: - single - deep_ensemble - bootstrap + - kfold - snapshot_ensemble description: Strategy registry key accepted by ``StrategyRegistry.get``. description: Strategy schemas backed by ``StrategyRegistry`` classes. diff --git a/xanesnet/strategies/__init__.py b/xanesnet/strategies/__init__.py index 39855667..86dbd8cc 100644 --- a/xanesnet/strategies/__init__.py +++ b/xanesnet/strategies/__init__.py @@ -23,6 +23,7 @@ from .base import Strategy from .bootstrap import Bootstrap from .deep_ensemble import DeepEnsemble +from .kfold import KFold from .registry import StrategyRegistry from .single import Single from .snapshot_ensemble import SnapshotEnsemble @@ -31,6 +32,7 @@ "Strategy", "Bootstrap", "DeepEnsemble", + "KFold", "SnapshotEnsemble", "Single", "StrategyRegistry", diff --git a/xanesnet/strategies/kfold.py b/xanesnet/strategies/kfold.py new file mode 100644 index 00000000..97013d66 --- /dev/null +++ b/xanesnet/strategies/kfold.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: GPL-3.0-or-later +# +# XANESNET +# +# Authors: Hendrik Junkawitsch, Tom J. Penfold, Tom W. Pope, C. D. Rankine, B. Li +# +# This program is free software: you can redistribute it and/or modify it under the terms of the +# GNU General Public License as published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without +# even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU +# General Public License for more details. +# +# You should have received a copy of the GNU General Public License along with this program. +# If not, see . +# +# Citations: +# ... + +"""K-fold cross-validation training and inference strategy for XANESNET.""" + +import copy +import logging +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import Subset + +from xanesnet.datasets import Dataset +from xanesnet.encodings import SpectraEncoding +from xanesnet.models import Model, ModelRegistry +from xanesnet.runners.inferencers import InferencerRegistry +from xanesnet.runners.trainers import TrainerRegistry +from xanesnet.serialization.config import Config +from xanesnet.serialization.tensorboard import tb_logger + +from .base import Strategy +from .registry import StrategyRegistry + + +@StrategyRegistry.register("kfold") +class KFold(Strategy): + """Repeated k-fold cross-validation strategy returning the best fold model. + + The strategy trains one model per fold on a shuffled partition of the full + dataset. Each fold uses the holdout partition as validation during training. + After all folds complete, the model with the lowest validation score is + returned for inference. + + Args: + strategy_type: Registry key identifying this strategy type. + dataset: Dataset used for training or inference. + model_config: Configuration for the model. + encoding: Composed spectra encoding forwarded to the trainers and + inferencer. + weight_init: Weight initialization scheme name. + weight_init_params: Additional weight-initializer parameters. + bias_init: Bias initialization scheme name. + n_splits: Number of folds per repeat. + n_repeats: Number of times to repeat the k-fold split. + seed: Random seed used to shuffle samples before splitting. + checkpoint_dir: Directory for checkpoints, or ``None``. + checkpoint_interval: Epoch interval between checkpoints, or ``None``. + tensorboard_dir: Directory for TensorBoard event files, or ``None``. + trainer_config: Trainer configuration for training mode. + inferencer_config: Inferencer configuration for inference mode. + """ + + def __init__( + self, + strategy_type: str, + dataset: Dataset, + model_config: Config, + encoding: SpectraEncoding, + weight_init: str, + weight_init_params: Config, + bias_init: str, + checkpoint_dir: str | Path | None, + checkpoint_interval: int | None, + tensorboard_dir: str | Path | None, + n_splits: int = 3, + n_repeats: int = 1, + seed: int | None = None, + trainer_config: Config | None = None, + inferencer_config: Config | None = None, + ) -> None: + """Initialize the k-fold cross-validation strategy.""" + super().__init__( + strategy_type, + dataset, + model_config, + encoding, + weight_init, + weight_init_params, + bias_init, + checkpoint_dir, + checkpoint_interval, + tensorboard_dir, + trainer_config, + inferencer_config, + ) + + if n_splits < 2: + raise ValueError(f"n_splits must be at least 2, got {n_splits}.") + if n_repeats < 1: + raise ValueError(f"n_repeats must be at least 1, got {n_repeats}.") + if len(self.dataset) < n_splits: + raise ValueError( + f"Dataset has {len(self.dataset)} samples, but k-fold requires at least {n_splits}." + ) + + self.n_splits = n_splits + self.n_repeats = n_repeats + self.seed = seed if seed is not None else np.random.default_rng().integers(0, 1000) + self._rng = np.random.default_rng(self.seed) + + self.model: Model | None = None + self.trainer: Any | None = None + self.inferencer: Any | None = None + self._device: str | torch.device | None = None + + def _iter_kfold_splits(self) -> Iterator[tuple[list[int], list[int]]]: + """Yield train and validation index lists for each fold. + + Yields: + Tuples of ``(train_indices, valid_indices)`` for one fold. + """ + n_samples = len(self.dataset) + indices = np.arange(n_samples) + + for _ in range(self.n_repeats): + shuffled = self._rng.permutation(indices) + fold_sizes = np.full(self.n_splits, n_samples // self.n_splits, dtype=int) + fold_sizes[: n_samples % self.n_splits] += 1 + + current = 0 + for fold_size in fold_sizes: + test_indices = shuffled[current : current + fold_size] + train_indices = np.concatenate([shuffled[:current], shuffled[current + fold_size :]]) + current += fold_size + yield train_indices.tolist(), test_indices.tolist() + + def _fold_dataset(self, train_indices: list[int], valid_indices: list[int]) -> Dataset: + """Return a dataset copy configured for one k-fold split. + + Args: + train_indices: Training indices for this fold. + valid_indices: Validation indices for this fold. + + Returns: + A shallow copy of ``self.dataset`` with train and validation + subsets set to the provided index lists. + """ + dataset_fold = copy.copy(self.dataset) + dataset_fold._subsets = [ + Subset(self.dataset, train_indices), + Subset(self.dataset, valid_indices), + ] + return dataset_fold + + def setup_models(self) -> None: + """Instantiate a template model from ``model_config`` for signatures.""" + model_type = self.model_config.get_str("model_type") + logging.info(f"Initializing k-fold model template: {model_type}") + self.model = ModelRegistry.create(model_type, **self.model_config.as_kwargs()) + + def init_model_weights(self) -> None: + """Apply weight and bias initialization to the template model.""" + if self.model is None: + raise ValueError("Cannot initialize model weights because the model is not initialized.") + + logging.info(f"Initializing weights with '{self.weight_init}' and bias with '{self.bias_init}'") + self.model.init_weights(self.weight_init, self.bias_init, **self.weight_init_params.as_kwargs()) + + def set_state_dicts(self, state_dicts: list[dict]) -> None: + """Load model weights from the first entry of ``state_dicts``. + + Args: + state_dicts: List of state dictionaries; only the first entry is + used for the selected k-fold model. + + Raises: + ValueError: If ``setup_models`` has not been called. + """ + if self.model is None: + raise ValueError("Cannot load state dicts because the model is not initialized.") + + self.model.load_state_dict(state_dicts[0]) + + def setup_trainers(self, device: str | torch.device) -> None: + """Store the training device; trainers are created per fold at runtime. + + Must be called after ``setup_models`` and ``setup_checkpointer``. + + Args: + device: The device on which training will be performed. + + Raises: + ValueError: If the model, trainer config, or checkpointer are not initialized. + """ + if self.model is None: + raise ValueError("Cannot setup trainers because the model is not initialized.") + if self.trainer_config is None: + raise ValueError("Can not setup trainers because there is no trainer config.") + if self.checkpointer is None: + raise ValueError("Can not setup trainers because checkpointer is not instantiated.") + + self._device = device + self.trainer = None + + def run_training(self) -> list[Model]: + """Train one model per fold and return the best-scoring model. + + Must be called after ``setup_trainers``. + + Returns: + A single-element list containing the fold model with the lowest + validation score. + + Raises: + ValueError: If setup steps were not completed or no fold produced + a usable validation score. + """ + if self.model is None: + raise ValueError("Cannot run training because the model is not initialized.") + if self.trainer_config is None: + raise ValueError("Cannot run training because there is no trainer config.") + if self.checkpointer is None: + raise ValueError("Cannot run training because checkpointer is not instantiated.") + if self._device is None: + raise ValueError("Cannot run training because trainers are not initialized.") + + super().run_training() + + model_type = self.model_config.get_str("model_type") + model_cls = ModelRegistry.get(model_type) + trainer_type = self.trainer_config.get_str("trainer_type") + trainer_cls = TrainerRegistry.get(trainer_type) + + best_model: Model | None = None + best_score = float("inf") + valid_scores: list[float] = [] + n_folds = self.n_splits * self.n_repeats + + for fold_idx, (train_indices, valid_indices) in enumerate(self._iter_kfold_splits()): + logging.info(f"Training k-fold model {fold_idx + 1}/{n_folds}.") + self.checkpointer.new_model() + + model = model_cls(**self.model_config.as_kwargs()) + model.init_weights(self.weight_init, self.bias_init, **self.weight_init_params.as_kwargs()) + dataset_fold = self._fold_dataset(train_indices, valid_indices) + trainer = trainer_cls( + **self.trainer_config.as_kwargs(), + dataset=dataset_fold, + model=model, + device=self._device, + checkpointer=self.checkpointer, + encoding=self.encoding, + ) + + try: + if self.tensorboard_dir is not None: + tb_logger.new_run(Path(self.tensorboard_dir) / f"fold_{fold_idx}") + + score = trainer.train() + finally: + tb_logger.close() + + if score is None: + logging.warning(f"Fold {fold_idx + 1} did not produce a validation score and will be skipped.") + model.to(torch.device("cpu")) + continue + + valid_scores.append(score) + + logging.info(f"Fold {fold_idx + 1} validation score: {score:.6f}") + if score < best_score: + logging.info(f"New best k-fold model found with validation score: {score:.6f}") + best_score = score + best_model = copy.deepcopy(model) + + model.to(torch.device("cpu")) + + if best_model is None: + raise ValueError("K-fold training did not produce a model with a validation score.") + + logging.info("K-fold cross-validation finished.") + if valid_scores: + logging.info( + f"Average validation score: {np.mean(valid_scores):.6f} +/- {np.std(valid_scores):.6f}" + ) + + self.model = best_model + return [self.model] + + def setup_inferencers(self, device: str | torch.device) -> None: + """Instantiate an inferencer for the selected k-fold model. + + Must be called after ``setup_models``. + + Args: + device: The device on which inference will be performed. + + Raises: + ValueError: If the model or inferencer config are not initialized. + """ + if self.model is None: + raise ValueError("Can not setup inferencers because the model is not initialized.") + if self.inferencer_config is None: + raise ValueError("Can not setup inferencers because there is no inferencer config.") + + inferencer_type = self.inferencer_config.get_str("inferencer_type") + logging.info(f"Initializing inferencer: {inferencer_type}") + + inferencer = InferencerRegistry.create( + inferencer_type, + **self.inferencer_config.as_kwargs(), + dataset=self.dataset, + model=self.model, + device=device, + encoding=self.encoding, + ) + + self.inferencer = inferencer + + def run_inference(self, predictions_save_path: str | Path | None) -> None: + """Run inference with the selected k-fold model. + + Args: + predictions_save_path: Directory in which to write prediction + output, or ``None`` to skip saving. + + Raises: + ValueError: If ``setup_inferencers`` has not been called. + """ + if self.inferencer is None: + raise ValueError("Cannot run inference because the Inferencer is not initialized.") + + super().run_inference(predictions_save_path) + + self.inferencer.infer(predictions_save_path) + + @property + def model_signature(self) -> Config: + """Return the model architecture signature. + + Returns: + A ``Config`` representing the model signature. + + Raises: + ValueError: If ``setup_models`` has not been called. + """ + if self.model is None: + raise ValueError("Model is not initialized. Cannot retrieve signature.") + + return self.model.signature + + @property + def signature(self) -> Config: + """Return the strategy configuration as a ``Config``. + + Returns: + A ``Config`` capturing the strategy configuration. + """ + signature = super().signature + signature.update_with_dict( + { + "n_splits": self.n_splits, + "n_repeats": self.n_repeats, + "seed": self.seed, + } + ) + return signature