"""Training engine and trained-surrogate containers.
This module contains the high-level training entrypoint :func:`train_surrogate` and
the thin wrappers (:class:`SurrogateModel`, :class:`TrainedModel`) used throughout
the package.
"""
from __future__ import annotations
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
import time
from typing import Any, Dict, Mapping, Tuple
import numpy as np
import torch
from cann.config.model_config import ModelConfig
from cann.data.datasets import ParamSurfaceDataset
from cann.engine.training_config import TrainingConfig
from cann.surrogate.mlp import build_mlp_network
def _scheduled_learning_rate(training_config: TrainingConfig, epoch: int) -> float:
"""Return the learning rate to use for a zero-based epoch."""
if training_config.lr_decay_every_epochs is None:
return training_config.learning_rate
decay_steps = epoch // training_config.lr_decay_every_epochs
learning_rate = training_config.learning_rate * (
training_config.lr_decay_factor**decay_steps
)
return max(training_config.min_learning_rate, learning_rate)
def _set_learning_rate(optimizer: torch.optim.Optimizer, learning_rate: float) -> None:
for param_group in optimizer.param_groups:
param_group["lr"] = learning_rate
def _clone_state_dict(network: torch.nn.Module) -> Dict[str, torch.Tensor]:
return {
name: value.detach().cpu().clone()
for name, value in network.state_dict().items()
}
def _atomic_torch_save(payload: Any, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
temporary_path = path.with_suffix(f"{path.suffix}.tmp")
torch.save(payload, temporary_path)
temporary_path.replace(path)
def _evaluate_tensors(
network: torch.nn.Module,
X: torch.Tensor,
y: torch.Tensor,
*,
batch_size: int,
) -> Dict[str, float]:
if X.shape[0] == 0:
raise ValueError("Cannot evaluate an empty dataset.")
network.eval()
squared_error = 0.0
absolute_error = 0.0
max_absolute_error = 0.0
num_values = 0
with torch.no_grad():
for start_idx in range(0, X.shape[0], batch_size):
outputs = network(X[start_idx : start_idx + batch_size])
diff = outputs - y[start_idx : start_idx + batch_size]
squared_error += torch.sum(diff * diff).item()
absolute_error += torch.sum(torch.abs(diff)).item()
max_absolute_error = max(
max_absolute_error, torch.max(torch.abs(diff)).item()
)
num_values += diff.numel()
mse = squared_error / num_values
return {
"mse": mse,
"rmse": mse**0.5,
"mae": absolute_error / num_values,
"max_abs_error": max_absolute_error,
}
[docs]
def evaluate_surrogate_dataset(
surrogate: SurrogateModel,
dataset: ParamSurfaceDataset,
*,
batch_size: int = 8192,
device: str | torch.device = "cpu",
) -> Dict[str, float]:
"""Evaluate a trained surrogate against an observed dataset."""
if dataset._observed_t is None:
raise ValueError("Evaluation dataset must have observed values.")
device = torch.device(device)
surrogate.network.to(device)
X = torch.cat([dataset._params_t, dataset._market_t], dim=1).to(device)
y = dataset._observed_t.to(device)
return _evaluate_tensors(surrogate.network, X, y, batch_size=batch_size)
def _checkpoint_payload(
*,
network: torch.nn.Module,
optimizer: torch.optim.Optimizer,
training_config: TrainingConfig,
completed_epochs: int,
best_model_state: Mapping[str, torch.Tensor] | None,
best_validation_mse: float | None,
best_epoch: int | None,
epoch_history: list[dict[str, float | int | None]],
metadata: Mapping[str, Any] | None,
) -> Dict[str, Any]:
payload = {
"format_version": 1,
"completed_epochs": completed_epochs,
"model_state_dict": _clone_state_dict(network),
"optimizer_state_dict": optimizer.state_dict(),
"training_config": asdict(training_config),
"best_model_state_dict": best_model_state,
"best_validation_mse": best_validation_mse,
"best_epoch": best_epoch,
"epoch_history": epoch_history,
"torch_rng_state": torch.get_rng_state(),
"metadata": dict(metadata or {}),
}
if torch.cuda.is_available():
payload["cuda_rng_state_all"] = torch.cuda.get_rng_state_all()
return payload
def _validate_resume_config(
saved_config: Mapping[str, Any], training_config: TrainingConfig
) -> None:
current_config = asdict(training_config)
saved_epochs = saved_config.get("epochs")
saved_without_epochs = dict(saved_config)
current_without_epochs = dict(current_config)
saved_without_epochs.pop("epochs", None)
current_without_epochs.pop("epochs", None)
if saved_without_epochs != current_without_epochs:
raise ValueError(
"Resume checkpoint training configuration does not match the current "
"training configuration."
)
if saved_epochs is not None and training_config.epochs < int(saved_epochs):
raise ValueError(
"Current training epoch target is smaller than the checkpoint's "
f"configured target: {training_config.epochs} < {saved_epochs}."
)
def _restore_training_checkpoint(
*,
checkpoint_path: str | Path,
network: torch.nn.Module,
optimizer: torch.optim.Optimizer,
training_config: TrainingConfig,
device: torch.device,
metadata: Mapping[str, Any] | None,
) -> tuple[
int,
Mapping[str, torch.Tensor] | None,
float | None,
int | None,
list[dict[str, float | int | None]],
]:
path = Path(checkpoint_path)
if not path.is_file():
raise FileNotFoundError(f"Resume checkpoint not found: {path}")
payload = torch.load(path, map_location=device, weights_only=False)
if payload.get("format_version") != 1:
raise ValueError(f"Unsupported checkpoint format in {path}")
_validate_resume_config(payload["training_config"], training_config)
expected_metadata = dict(metadata or {})
saved_metadata = payload.get("metadata", {})
for key, expected_value in expected_metadata.items():
if saved_metadata.get(key) != expected_value:
raise ValueError(
f"Resume checkpoint metadata mismatch for '{key}': "
f"{saved_metadata.get(key)!r} != {expected_value!r}"
)
completed_epochs = int(payload["completed_epochs"])
if completed_epochs > training_config.epochs:
raise ValueError(
f"Checkpoint has completed {completed_epochs} epochs, but the current "
f"target is only {training_config.epochs}."
)
network.load_state_dict(payload["model_state_dict"])
optimizer.load_state_dict(payload["optimizer_state_dict"])
torch.set_rng_state(payload["torch_rng_state"].cpu())
if device.type == "cuda" and "cuda_rng_state_all" in payload:
torch.cuda.set_rng_state_all(payload["cuda_rng_state_all"])
return (
completed_epochs,
payload.get("best_model_state_dict"),
payload.get("best_validation_mse"),
payload.get("best_epoch"),
payload.get("epoch_history", []),
)
[docs]
@dataclass
class SurrogateModel:
"""
Wrapper around the neural network surrogate, scalers and metadata.
Fields:
``network`` (PyTorch module), ``model_cfg`` (feature/parameter layout), and optional
``version`` string.
"""
network: torch.nn.Module
model_cfg: ModelConfig
version: str = "0.0.1"
[docs]
def to(self, device: torch.device | str) -> SurrogateModel:
"""
Move the surrogate model to the specified device.
Args:
device (torch.device | str): Target device (e.g., "cpu", "cuda", "mps").
Returns:
SurrogateModel: Self, moved to the requested device.
"""
self.network.to(device)
return self
[docs]
def eval(self) -> SurrogateModel:
"""
Set the surrogate model to evaluation mode.
Returns:
SurrogateModel: Self, in evaluation mode.
"""
self.network.eval()
return self
@torch.no_grad()
def predict(
self,
parameters: np.ndarray,
market_inputs: np.ndarray,
batch_size: int = 8192,
device: str | torch.device = "cpu",
) -> np.ndarray:
"""
Convenience prediction helper for calibration / sanity checks.
Args:
parameters (np.ndarray): Array of model parameters with shape (N, n).
market_inputs (np.ndarray): Array of market input features with shape (N, m).
batch_size (int): Batch size for prediction. Defaults to 8192.
device (torch.device | str): Device to run prediction on. Defaults to "cpu".
Returns:
np.ndarray: Predicted outputs with shape (N, output_dim).
"""
self.network.eval()
self.network.to(device)
device = torch.device(device)
# Use DataLoader for efficient batching
dataset = ParamSurfaceDataset(
parameters=parameters,
market_inputs=market_inputs,
model_cfg=self.model_cfg,
)
loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=False
)
preds = []
for b in loader:
batch = torch.cat([b["params"], b["market"]], dim=1).to(device)
y = self.network(batch)
preds.append(y.cpu().numpy())
return np.concatenate(preds, axis=0)
def _get_device(
self, training_config: TrainingConfig, verbose: bool
) -> torch.device:
if training_config.device == "cuda" and torch.cuda.is_available():
if verbose:
print("Using CUDA for training.")
return torch.device("cuda")
elif training_config.device == "mps" and torch.backends.mps.is_available():
if verbose:
print("Using MPS for training.")
return torch.device("mps")
else:
if verbose:
print("Using CPU for training.")
return torch.device("cpu")
[docs]
def train(
self,
train_dataset: ParamSurfaceDataset,
validation_dataset: ParamSurfaceDataset | None,
training_config: TrainingConfig,
verbose: bool = True,
checkpoint_dir: str | Path | None = None,
checkpoint_every_epochs: int | None = None,
resume_checkpoint: str | Path | None = None,
checkpoint_metadata: Mapping[str, Any] | None = None,
) -> TrainedModel:
"""
Train the surrogate model on the provided training dataset.
Args:
train_dataset (ParamSurfaceDataset): Dataset for training.
validation_dataset (ParamSurfaceDataset | None): Dataset for validation.
training_config (TrainingConfig): Configuration for training.
verbose (bool): Whether to print training progress. Defaults to True.
checkpoint_dir (str | Path | None): Optional directory for resumable checkpoints.
checkpoint_every_epochs (int | None): Save a checkpoint after this many epochs.
resume_checkpoint (str | Path | None): Optional checkpoint to resume from.
checkpoint_metadata (Mapping[str, Any] | None): Provenance stored in checkpoints.
Returns:
TrainedModel: A trained model wrapper containing this surrogate and training metadata.
Raises:
ValueError: If the provided datasets do not contain observed targets.
"""
if training_config.random_seed is not None:
torch.manual_seed(training_config.random_seed)
if checkpoint_every_epochs is not None and checkpoint_every_epochs <= 0:
raise ValueError(
"checkpoint_every_epochs must be positive when provided, "
f"got {checkpoint_every_epochs}"
)
if checkpoint_dir is not None and checkpoint_every_epochs is None:
raise ValueError(
"checkpoint_every_epochs must be provided when checkpoint_dir is set."
)
if checkpoint_dir is None and checkpoint_every_epochs is not None:
raise ValueError(
"checkpoint_dir must be provided when checkpoint_every_epochs is set."
)
if train_dataset._observed_t is None:
raise ValueError("Training dataset must have observed values for training.")
if validation_dataset is not None and validation_dataset._observed_t is None:
raise ValueError(
"Validation dataset must have observed values for validation."
)
device = self._get_device(training_config, verbose)
self.network.to(device)
# 2. OPTIMIZATION: Move ENTIRE dataset to GPU/Device upfront.
# Accessing protected members _params_t directly for speed.
X_train = torch.cat(
[train_dataset._params_t, train_dataset._market_t], dim=1
).to(device)
y_train = train_dataset._observed_t.to(device)
N_train = X_train.shape[0]
if N_train == 0:
raise ValueError("Training split is empty.")
X_val, y_val = None, None
if validation_dataset is not None:
X_val = torch.cat(
[validation_dataset._params_t, validation_dataset._market_t], dim=1
).to(device)
y_val = validation_dataset._observed_t.to(device) # type: ignore
if training_config.optimizer.lower() == "adam":
optimizer = torch.optim.Adam(
self.network.parameters(), lr=training_config.learning_rate
)
elif training_config.optimizer.lower() == "sgd":
optimizer = torch.optim.SGD(
self.network.parameters(), lr=training_config.learning_rate
)
else:
raise ValueError(f"Unsupported optimizer: {training_config.optimizer}")
if training_config.loss.lower() != "mse":
raise ValueError(f"Unsupported loss: {training_config.loss}")
loss_fn = torch.nn.MSELoss()
batch_size = training_config.batch_size
epoch_times = []
epoch_history: list[dict[str, float | int | None]] = []
start_epoch = 0
best_model_state: Mapping[str, torch.Tensor] | None = None
best_validation_mse: float | None = None
best_epoch: int | None = None
if resume_checkpoint is not None:
(
start_epoch,
best_model_state,
best_validation_mse,
best_epoch,
epoch_history,
) = _restore_training_checkpoint(
checkpoint_path=resume_checkpoint,
network=self.network,
optimizer=optimizer,
training_config=training_config,
device=device,
metadata=checkpoint_metadata,
)
if verbose:
print(f"Resuming training from epoch {start_epoch}.")
# 3. Training Loop
epoch_loss = float("nan")
for epoch in range(start_epoch, training_config.epochs):
start_time = time.time()
learning_rate = _scheduled_learning_rate(training_config, epoch)
_set_learning_rate(optimizer, learning_rate)
# Shuffle indices on the GPU
if training_config.shuffle:
indices = torch.randperm(N_train, device=device)
else:
indices = torch.arange(N_train, device=device)
self.network.train()
epoch_loss = 0.0
# Iterate via slicing (Vectorized)
for start_idx in range(0, N_train, batch_size):
idx = indices[start_idx : start_idx + batch_size]
# Direct slice - no dictionary overhead
batch_X = X_train[idx]
batch_y = y_train[idx]
optimizer.zero_grad()
outputs = self.network(batch_X)
loss = loss_fn(outputs, batch_y)
loss.backward()
optimizer.step()
epoch_loss += loss.item() * len(idx)
# Validation
train_loss = epoch_loss / N_train
validation_mse = None
should_validate = X_val is not None and (
epoch == 0
or (epoch + 1) % training_config.validation_every_epochs == 0
or epoch + 1 == training_config.epochs
)
if should_validate and X_val is not None and y_val is not None:
validation_mse = _evaluate_tensors(
self.network, X_val, y_val, batch_size=batch_size
)["mse"]
if (
best_validation_mse is None
or validation_mse < best_validation_mse
):
best_validation_mse = validation_mse
best_epoch = epoch + 1
best_model_state = _clone_state_dict(self.network)
end_time = time.time()
epoch_duration = end_time - start_time
epoch_times.append(epoch_duration)
epoch_history.append(
{
"epoch": epoch + 1,
"learning_rate": learning_rate,
"train_mse": train_loss,
"validation_mse": validation_mse,
"duration_seconds": epoch_duration,
}
)
should_report = (
epoch == 0
or (epoch + 1) % 100 == 0
or epoch + 1 == training_config.epochs
)
if verbose and should_report:
val_text = (
f"{validation_mse:.6g}"
if validation_mse is not None
else "not evaluated"
)
print(
f"Epoch [{epoch + 1}/{training_config.epochs}], "
f"Train MSE: {train_loss:.6g}, "
f"Val MSE: {val_text}, "
f"LR: {learning_rate:.6g}, "
f"Time per epoch: {epoch_duration:.4f}s"
)
completed_epochs = epoch + 1
should_checkpoint = (
checkpoint_dir is not None
and checkpoint_every_epochs is not None
and (
completed_epochs % checkpoint_every_epochs == 0
or completed_epochs == training_config.epochs
)
)
if should_checkpoint:
payload = _checkpoint_payload(
network=self.network,
optimizer=optimizer,
training_config=training_config,
completed_epochs=completed_epochs,
best_model_state=best_model_state,
best_validation_mse=best_validation_mse,
best_epoch=best_epoch,
epoch_history=epoch_history,
metadata=checkpoint_metadata,
)
checkpoint_path = (
Path(checkpoint_dir)
/ f"checkpoint_epoch_{completed_epochs:05d}.pth"
)
_atomic_torch_save(payload, checkpoint_path)
_atomic_torch_save(
payload, Path(checkpoint_dir) / "checkpoint_latest.pth"
)
if verbose:
print(f"Saved resumable checkpoint: {checkpoint_path}")
if verbose and epoch_times:
avg_time = sum(epoch_times) / len(epoch_times)
print(f"Average time per epoch: {avg_time:.4f}s")
if best_model_state is not None:
self.network.load_state_dict(best_model_state)
selected_train_metrics = _evaluate_tensors(
self.network, X_train, y_train, batch_size=batch_size
)
selected_validation_metrics = (
_evaluate_tensors(self.network, X_val, y_val, batch_size=batch_size)
if X_val is not None and y_val is not None
else None
)
final_epoch_train_mse = (
epoch_loss / N_train
if epoch_times
else float(epoch_history[-1]["train_mse"])
)
return TrainedModel(
surrogate=self,
training_config=training_config,
model_config=self.model_cfg,
training_metrics={
"completed_epochs": training_config.epochs,
"final_epoch_train_mse": final_epoch_train_mse,
"best_epoch": best_epoch,
"best_validation_mse": best_validation_mse,
"selected_train": selected_train_metrics,
"selected_validation": selected_validation_metrics,
"epoch_history": epoch_history,
},
)
[docs]
@dataclass
class TrainedModel:
"""
Container for a trained surrogate model along with training metadata.
Fields:
``surrogate`` plus associated configs and training metadata.
"""
surrogate: SurrogateModel
training_config: TrainingConfig
model_config: ModelConfig
training_metrics: Dict[str, Any]
path: str | None = None
training_date: datetime = field(default_factory=datetime.now)
[docs]
def train_surrogate(
model_config: ModelConfig,
parameters: np.ndarray,
market_inputs: np.ndarray,
observed: np.ndarray,
verbose: bool | None = True,
training_config: TrainingConfig | None = None,
checkpoint_dir: str | Path | None = None,
checkpoint_every_epochs: int | None = None,
resume_checkpoint: str | Path | None = None,
checkpoint_metadata: Mapping[str, Any] | None = None,
) -> Tuple[TrainedModel, ParamSurfaceDataset | None]:
"""
Train a neural network surrogate model on the provided dataset.
The `parameters` and `market_inputs` should be numpy arrays with columns in the order
of the `parameters` and `market_inputs` defined by the `model_config` provided.
This is the high-level entry point used by the public API.
Implementation will perform:
- dataset split
- surrogate construction
- training loop
- metric computation
Args:
model_config (ModelConfig): Configuration for the underlying pricing model.
parameters (np.ndarray): Array of model parameters with shape (N, n).
market_inputs (np.ndarray): Array of market input features with shape (N, m).
observed (np.ndarray): Array of observed market values with shape (N, ) or (N, 1).
verbose (bool): Whether to print training progress. Defaults to True.
training_config (TrainingConfig | None): Training hyperparameters. If not provided,
:attr:`ModelConfig.train_defaults` is used.
checkpoint_dir (str | Path | None): Optional directory for resumable checkpoints.
checkpoint_every_epochs (int | None): Save a checkpoint after this many epochs.
resume_checkpoint (str | Path | None): Optional checkpoint to resume from.
checkpoint_metadata (Mapping[str, Any] | None): Provenance stored in checkpoints.
Returns:
Tuple[TrainedModel, ParamSurfaceDataset | None]: The trained model and (optionally)
the test split dataset (``None`` if the test fraction is 0).
Raises:
ValueError: If no training configuration is available.
"""
if training_config is None:
training_config = model_config.train_defaults
if verbose:
print(training_config)
if training_config is None:
raise ValueError(
"No training configuration provided. Please specify training_config or set model_config.train_defaults."
)
if training_config.random_seed is not None:
torch.manual_seed(training_config.random_seed)
# Split the data into train, validation, and test sets and then put into ParamSurfaceDataset
num_samples = parameters.shape[0]
train_frac, val_frac, test_frac = training_config.train_val_test_split
train_end = int(train_frac * num_samples)
val_end = train_end + int(val_frac * num_samples)
train_dataset = ParamSurfaceDataset(
parameters=parameters[:train_end],
market_inputs=market_inputs[:train_end],
observed=observed[:train_end],
model_cfg=model_config,
)
validation_dataset = None
if val_frac > 0:
validation_dataset = ParamSurfaceDataset(
parameters=parameters[train_end:val_end],
market_inputs=market_inputs[train_end:val_end],
observed=observed[train_end:val_end],
model_cfg=model_config,
)
test_dataset = None
if test_frac > 0:
test_dataset = ParamSurfaceDataset(
parameters=parameters[val_end:],
market_inputs=market_inputs[val_end:],
observed=observed[val_end:],
model_cfg=model_config,
)
# Build the surrogate network
network = build_mlp_network(model_config, training_config)
surrogate = SurrogateModel(network=network, model_cfg=model_config)
trained_model = surrogate.train(
train_dataset,
validation_dataset,
training_config,
verbose if verbose is not None else True,
checkpoint_dir=checkpoint_dir,
checkpoint_every_epochs=checkpoint_every_epochs,
resume_checkpoint=resume_checkpoint,
checkpoint_metadata=checkpoint_metadata,
)
return trained_model, test_dataset