Source code for cann.engine.training_config

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Tuple, List, Literal




@dataclass
class TrainingConfig:
    """
    Hyperparameters and settings for training the surrogate model.

    Commonly adjusted fields are ``epochs``, ``batch_size``, ``learning_rate``,
    ``hidden_layers``, ``activation``, and ``train_val_test_split``.
    """

    batch_size: int = 32768
    epochs: int = 4800
    dropout_rate: float = 0.0
    batch_norm: bool = False
    learning_rate: float = 1e-3
    lr_decay_every_epochs: int | None = None
    lr_decay_factor: float = 0.5
    min_learning_rate: float = 0.0
    optimizer: Literal["adam", "sgd"] = "adam"
    loss: Literal["mse"] = "mse"
    validation_every_epochs: int = 100
    hidden_layers: List[int] = field(default_factory=lambda: [200, 200, 200, 200])
    activation: Literal["relu", "tanh", "gelu"] = "relu"
    weight_init: Literal["default", "uniform"] = "default"
    weight_init_min: float = -0.05
    weight_init_max: float = 0.05
    bias_init: float | None = None
    train_val_test_split: Tuple[float, float, float] = (0.8, 0.1, 0.1)
    shuffle: bool = True
    early_stopping: bool = False
    device: Literal["cpu", "cuda"] = "cuda"
    random_seed: int | None = 42

    def __post_init__(self) -> None:
        s = sum(self.train_val_test_split)
        if not (0.99 <= s <= 1.01):
            raise ValueError(f"train_val_test_split must sum to 1.0, got {s}")
        if self.batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {self.batch_size}")
        if self.epochs <= 0:
            raise ValueError(f"epochs must be positive, got {self.epochs}")
        if self.learning_rate <= 0:
            raise ValueError(
                f"learning_rate must be positive, got {self.learning_rate}"
            )
        if self.lr_decay_every_epochs is not None and self.lr_decay_every_epochs <= 0:
            raise ValueError(
                "lr_decay_every_epochs must be positive when provided, "
                f"got {self.lr_decay_every_epochs}"
            )
        if not (0 < self.lr_decay_factor <= 1):
            raise ValueError(
                f"lr_decay_factor must be in (0, 1], got {self.lr_decay_factor}"
            )
        if self.min_learning_rate < 0:
            raise ValueError(
                f"min_learning_rate must be non-negative, got {self.min_learning_rate}"
            )
        if self.min_learning_rate > self.learning_rate:
            raise ValueError(
                "min_learning_rate must not exceed learning_rate, "
                f"got {self.min_learning_rate} > {self.learning_rate}"
            )
        if self.validation_every_epochs <= 0:
            raise ValueError(
                "validation_every_epochs must be positive, "
                f"got {self.validation_every_epochs}"
            )
        if not self.hidden_layers or any(width <= 0 for width in self.hidden_layers):
            raise ValueError(
                f"hidden_layers must contain positive widths, got {self.hidden_layers}"
            )
        if self.weight_init not in ("default", "uniform"):
            raise ValueError(f"Unsupported weight_init: {self.weight_init}")
        if self.weight_init_min >= self.weight_init_max:
            raise ValueError(
                "weight_init_min must be smaller than weight_init_max, "
                f"got {self.weight_init_min} >= {self.weight_init_max}"
            )



    @classmethod
    def from_dict(cls, cfg: dict) -> TrainingConfig:
        """
        Create a TrainingConfig instance from a dictionary.

        Args:
            cfg (dict): Dictionary containing configuration parameters.

        Returns:
            TrainingConfig: Parsed training configuration.

        Raises:
            TypeError: If the dictionary contains keys not accepted by TrainingConfig.
        """
        return cls(**cfg)