marn 0.1.0__py3-none-any.whl

marn/__init__.py

@@ -0,0 +1,154 @@
+"""Train PyTorch models through low-dimensional parameter mappings."""
+
+from marn.generators import (
+    GroupedGenerator,
+    LayerGenerator,
+    LazyLayerwiseGenerator,
+    LayerwiseGenerator,
+    ParameterGenerator,
+    SingleVectorGenerator,
+    FineTuningGenerator,
+    LRDGenerator,
+)
+from marn.distributed import (
+    benchmark_strategies,
+    cleanup_ddp,
+    is_ddp_available,
+    profile_peak_memory,
+    setup_ddp,
+    wrap_ddp,
+)
+from marn.models import (
+    ForwardResult,
+    MappingModel,
+    TargetModel,
+    UnsupportedTargetModelError,
+)
+from marn.losses import (
+    AlignmentLoss,
+    BaseLoss,
+    ClassificationLoss,
+    LossOutput,
+    MappingLoss,
+    RegressionLoss,
+    SmoothnessLoss,
+    StabilityLoss,
+    TaskLoss,
+    TrainingContext,
+)
+from marn.mappers import BaseMapper, MLPMapper, ResidualMLPMapper
+from marn.modulation import (
+    AdditiveModulation,
+    AffineModulation,
+    BaseModulation,
+    LowRankModulation,
+)
+from marn.runtime import ParameterEntry, ParameterSpec, ParameterTree
+from marn.strategies import (
+    GroupedStrategy,
+    LayerwiseStrategy,
+    SLVTStrategy,
+    FineTuningStrategy,
+    LRDStrategy,
+)
+from marn.config import (
+    GeneratorConfig,
+    LossConfig,
+    MapperConfig,
+    MappingConfig,
+    TaskLossConfig,
+    TrainerConfig,
+    load_config,
+)
+from marn.registry import (
+    GENERATOR_REGISTRY,
+    LOSS_REGISTRY,
+    MAPPER_REGISTRY,
+    MODULATION_REGISTRY,
+    Registry,
+)
+from marn.callbacks import Callback, EarlyStopping, MetricLogger
+from marn.trainers import (
+    BatchAdapter,
+    MappingBatchAdapter,
+    MappingTrainer,
+    TupleBatchAdapter,
+    LRFinderResult,
+)
+from marn.checkpoint import (
+    CheckpointCompatibilityError,
+    CheckpointSchema,
+    load_checkpoint,
+    save_checkpoint,
+)
+
+__all__ = [
+    "AdditiveModulation",
+    "AffineModulation",
+    "AlignmentLoss",
+    "BaseLoss",
+    "BaseMapper",
+    "BaseModulation",
+    "BatchAdapter",
+    "Callback",
+    "CheckpointCompatibilityError",
+    "CheckpointSchema",
+    "ClassificationLoss",
+    "EarlyStopping",
+    "ForwardResult",
+    "FineTuningGenerator",
+    "GENERATOR_REGISTRY",
+    "GeneratorConfig",
+    "GroupedGenerator",
+    "GroupedStrategy",
+    "LOSS_REGISTRY",
+    "LayerGenerator",
+    "LazyLayerwiseGenerator",
+    "LayerwiseGenerator",
+    "LayerwiseStrategy",
+    "LRDGenerator",
+    "LossConfig",
+    "LossOutput",
+    "LowRankModulation",
+    "MAPPER_REGISTRY",
+    "MODULATION_REGISTRY",
+    "MapperConfig",
+    "MappingBatchAdapter",
+    "MappingConfig",
+    "MappingLoss",
+    "MappingModel",
+    "MappingTrainer",
+    "MLPMapper",
+    "MetricLogger",
+    "ParameterEntry",
+    "ParameterGenerator",
+    "ParameterSpec",
+    "ParameterTree",
+    "RegressionLoss",
+    "Registry",
+    "ResidualMLPMapper",
+    "SLVTStrategy",
+    "SingleVectorGenerator",
+    "FineTuningStrategy",
+    "LRDStrategy",
+    "LRFinderResult",
+    "SmoothnessLoss",
+    "StabilityLoss",
+    "TargetModel",
+    "TaskLoss",
+    "TaskLossConfig",
+    "TrainerConfig",
+    "TrainingContext",
+    "TupleBatchAdapter",
+    "UnsupportedTargetModelError",
+    "benchmark_strategies",
+    "cleanup_ddp",
+    "is_ddp_available",
+    "load_checkpoint",
+    "load_config",
+    "profile_peak_memory",
+    "save_checkpoint",
+    "setup_ddp",
+    "wrap_ddp",
+]
+__version__ = "0.1.0"

marn/callbacks/__init__.py

@@ -0,0 +1,7 @@
+"""Training lifecycle callbacks."""
+
+from marn.callbacks.base import Callback
+from marn.callbacks.early_stopping import EarlyStopping
+from marn.callbacks.logger import MetricLogger
+
+__all__ = ["Callback", "EarlyStopping", "MetricLogger"]

marn/callbacks/base.py

@@ -0,0 +1,51 @@
+"""Observer callback base class for training lifecycle hooks."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from marn.losses.outputs import LossOutput
+
+
+class Callback:
+    """Base class for training lifecycle callbacks.
+
+    All hooks are no-ops by default — subclass and override the ones you need.
+    Callbacks receive the trainer instance so they can inspect or modify state
+    (e.g. set ``trainer.should_stop = True`` for early termination).
+
+    Hook execution order within each event:
+
+    1. ``on_fit_start``
+    2. For each epoch:
+       a. ``on_epoch_start``
+       b. For each batch: ``on_batch_start`` → ``on_batch_end``
+       c. ``on_validation_start`` → ``on_validation_end``
+       d. ``on_epoch_end``
+    3. ``on_fit_end``
+    """
+
+    def on_fit_start(self, trainer: Any) -> None:
+        """Called once at the beginning of :meth:`fit`."""
+
+    def on_fit_end(self, trainer: Any) -> None:
+        """Called once at the end of :meth:`fit`."""
+
+    def on_epoch_start(self, trainer: Any, epoch: int) -> None:
+        """Called at the beginning of each epoch."""
+
+    def on_epoch_end(self, trainer: Any, epoch: int, metrics: dict[str, float]) -> None:
+        """Called at the end of each epoch with aggregated metrics."""
+
+    def on_batch_start(self, trainer: Any, batch_idx: int) -> None:
+        """Called before processing each batch."""
+
+    def on_batch_end(self, trainer: Any, batch_idx: int, loss_output: LossOutput) -> None:
+        """Called after processing each batch with the loss output."""
+
+    def on_validation_start(self, trainer: Any) -> None:
+        """Called before the validation loop."""
+
+    def on_validation_end(self, trainer: Any, metrics: dict[str, float]) -> None:
+        """Called after the validation loop with validation metrics."""

marn/callbacks/early_stopping.py

@@ -0,0 +1,60 @@
+"""Early stopping callback."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from marn.callbacks.base import Callback
+
+
+class EarlyStopping(Callback):
+    """Stop training when a monitored metric stops improving.
+
+    Args:
+        monitor: Name of the metric to monitor (e.g. ``"val_loss"``).
+        patience: Number of epochs with no improvement before stopping.
+        min_delta: Minimum change to qualify as an improvement.
+        mode: ``"min"`` to minimize the metric, ``"max"`` to maximize it.
+
+    Example::
+
+        trainer = MappingTrainer(
+            ...,
+            callbacks=[EarlyStopping(monitor="val_loss", patience=5)],
+        )
+    """
+
+    def __init__(
+        self,
+        monitor: str = "val_loss",
+        patience: int = 10,
+        min_delta: float = 0.0,
+        mode: str = "min",
+    ) -> None:
+        if mode not in ("min", "max"):
+            raise ValueError(f"mode must be 'min' or 'max', got {mode!r}")
+        self.monitor = monitor
+        self.patience = patience
+        self.min_delta = min_delta
+        self.mode = mode
+        self._best: float = math.inf if mode == "min" else -math.inf
+        self._wait: int = 0
+
+    def _is_improvement(self, current: float) -> bool:
+        if self.mode == "min":
+            return current < self._best - self.min_delta
+        return current > self._best + self.min_delta
+
+    def on_epoch_end(self, trainer: Any, epoch: int, metrics: dict[str, float]) -> None:
+        current = metrics.get(self.monitor)
+        if current is None:
+            return
+
+        if self._is_improvement(current):
+            self._best = current
+            self._wait = 0
+        else:
+            self._wait += 1
+            if self._wait >= self.patience:
+                trainer.should_stop = True

marn/callbacks/logger.py

@@ -0,0 +1,47 @@
+"""Structured metric logging callback."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from marn.callbacks.base import Callback
+
+if TYPE_CHECKING:
+    from marn.losses.outputs import LossOutput
+
+logger = logging.getLogger("marn.trainer")
+
+
+class MetricLogger(Callback):
+    """Log training and validation metrics using the standard ``logging`` module.
+
+    Args:
+        log_every_n_batches: Log batch metrics every N batches.
+
+    Example::
+
+        trainer = MappingTrainer(
+            ...,
+            callbacks=[MetricLogger(log_every_n_batches=50)],
+        )
+    """
+
+    def __init__(self, log_every_n_batches: int = 10) -> None:
+        self.log_every_n_batches = log_every_n_batches
+
+    def on_epoch_start(self, trainer: Any, epoch: int) -> None:
+        logger.info("Epoch %d/%d started", epoch + 1, trainer.config.max_epochs)
+
+    def on_batch_end(self, trainer: Any, batch_idx: int, loss_output: LossOutput) -> None:
+        if (batch_idx + 1) % self.log_every_n_batches == 0:
+            metrics_str = ", ".join(f"{k}={v:.4f}" for k, v in loss_output.metrics.items())
+            logger.info("  batch %d: %s", batch_idx + 1, metrics_str)
+
+    def on_epoch_end(self, trainer: Any, epoch: int, metrics: dict[str, float]) -> None:
+        metrics_str = ", ".join(f"{k}={v:.4f}" for k, v in metrics.items())
+        logger.info("Epoch %d/%d: %s", epoch + 1, trainer.config.max_epochs, metrics_str)
+
+    def on_validation_end(self, trainer: Any, metrics: dict[str, float]) -> None:
+        metrics_str = ", ".join(f"{k}={v:.4f}" for k, v in metrics.items())
+        logger.info("  validation: %s", metrics_str)

marn/checkpoint/__init__.py

@@ -0,0 +1,12 @@
+"""Versioned checkpoint save and load for mapping network training state."""
+
+from marn.checkpoint.save import save_checkpoint
+from marn.checkpoint.load import load_checkpoint, CheckpointCompatibilityError
+from marn.checkpoint.schema import CheckpointSchema
+
+__all__ = [
+    "CheckpointCompatibilityError",
+    "CheckpointSchema",
+    "load_checkpoint",
+    "save_checkpoint",
+]

marn/checkpoint/load.py

@@ -0,0 +1,174 @@
+"""Load a mapping network checkpoint from disk and restore model/trainer state."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from marn.checkpoint.schema import CURRENT_SCHEMA_VERSION, CheckpointSchema
+
+if TYPE_CHECKING:
+    from marn.models.mapping_model import MappingModel
+    from marn.trainers.trainer import MappingTrainer
+
+logger = logging.getLogger("marn.checkpoint")
+
+
+class CheckpointCompatibilityError(ValueError):
+    """Raised when a checkpoint's architecture is incompatible with the model.
+
+    Possible causes:
+
+    * Schema version from a future release is not understood by this version.
+    * Target ``ParameterSpec`` names or shapes differ (different architecture).
+    * Latent or buffer tensors do not match the current generator's keys.
+    """
+
+
+def _load_payload(path: Path, map_location: Any) -> dict[str, Any]:
+    raw = torch.load(path, map_location=map_location, weights_only=False)
+    if not isinstance(raw, dict):
+        raise CheckpointCompatibilityError(
+            f"Checkpoint at {path} is not a valid marn checkpoint "
+            f"(expected a dict, got {type(raw).__name__})."
+        )
+    return raw
+
+
+def _parse_schema(payload: dict[str, Any], path: Path) -> CheckpointSchema:
+    schema_version = payload.get("schema_version")
+    if schema_version is None:
+        raise CheckpointCompatibilityError(
+            f"Checkpoint at {path} is missing 'schema_version'. "
+            "It may have been saved by an incompatible tool."
+        )
+    if schema_version > CURRENT_SCHEMA_VERSION:
+        raise CheckpointCompatibilityError(
+            f"Checkpoint schema version {schema_version} is newer than the "
+            f"current package supports ({CURRENT_SCHEMA_VERSION}). "
+            "Upgrade marn to load this checkpoint."
+        )
+    return CheckpointSchema(
+        schema_version=schema_version,
+        package_version=payload.get("package_version", "unknown"),
+        latent_state=payload["latent_state"],
+        mapper_buffers=payload["mapper_buffers"],
+        parameter_spec_names=payload["parameter_spec_names"],
+        config=payload.get("config"),
+        trainer_state=payload.get("trainer_state"),
+        metadata=payload.get("metadata", {}),
+    )
+
+
+def _validate_spec(schema: CheckpointSchema, model: MappingModel, path: Path) -> None:
+    """Raise ``CheckpointCompatibilityError`` if the target spec has changed."""
+    spec = model.target.parameter_spec
+    current: list[tuple[str, list[int]]] = [(entry.name, list(entry.shape)) for entry in spec]
+    saved = [(name, shape) for name, shape in schema.parameter_spec_names]
+
+    if len(current) != len(saved):
+        raise CheckpointCompatibilityError(
+            f"Checkpoint at {path} has {len(saved)} target parameters but "
+            f"the model has {len(current)}.  The target architecture differs."
+        )
+
+    mismatches: list[str] = []
+    for (cur_name, cur_shape), (sav_name, sav_shape) in zip(current, saved):
+        if cur_name != sav_name:
+            mismatches.append(f"  name mismatch: checkpoint={sav_name!r} model={cur_name!r}")
+        elif cur_shape != sav_shape:
+            mismatches.append(
+                f"  shape mismatch for {cur_name!r}: checkpoint={sav_shape} model={cur_shape}"
+            )
+
+    if mismatches:
+        detail = "\n".join(mismatches)
+        raise CheckpointCompatibilityError(
+            f"Checkpoint at {path} is incompatible with this model:\n{detail}"
+        )
+
+
+def load_checkpoint(
+    path: str | Path,
+    model: MappingModel,
+    *,
+    trainer: MappingTrainer | None = None,
+    map_location: Any = None,
+    strict_latents: bool = True,
+) -> CheckpointSchema:
+    """Load a checkpoint and restore model (and optionally trainer) state.
+
+    The target's ``ParameterSpec`` is validated against the checkpoint before
+    any state is applied.  If the architecture differs a
+    ``CheckpointCompatibilityError`` is raised and the model is left unchanged.
+
+    Args:
+        path: Path to the ``.pt`` checkpoint file.
+        model: A ``MappingModel`` whose generator state will be restored.
+        trainer: Optional ``MappingTrainer`` whose optimizer, scheduler, epoch,
+            and GradScaler state will be restored when present in the checkpoint.
+        map_location: Passed to ``torch.load`` for device remapping.
+        strict_latents: When *True*, ``load_state_dict`` is called in strict
+            mode — every saved latent and buffer key must match the generator.
+            Set to *False* to allow partial loading (e.g. subsets of layers).
+
+    Returns:
+        The parsed ``CheckpointSchema`` so callers can inspect ``metadata``,
+        ``config``, or ``trainer_state`` independently.
+
+    Raises:
+        FileNotFoundError: If the checkpoint file does not exist.
+        CheckpointCompatibilityError: If schema version or architecture is
+            incompatible.
+
+    Example::
+
+        schema = load_checkpoint("checkpoints/epoch10.pt", model=model, trainer=trainer)
+        print(f"Resumed from epoch {schema.trainer_state['epoch']}")
+    """
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {path}")
+
+    payload = _load_payload(path, map_location)
+    schema = _parse_schema(payload, path)
+
+    # Validate architecture before touching any model state
+    _validate_spec(schema, model, path)
+
+    # Restore generator state (latents + buffers together)
+    combined_state: dict[str, Any] = {**schema.latent_state, **schema.mapper_buffers}
+    model.generator.load_state_dict(combined_state, strict=strict_latents)
+
+    logger.info(
+        "Checkpoint loaded from %s  (schema_version=%d  package=%s)",
+        path,
+        schema.schema_version,
+        schema.package_version,
+    )
+
+    # Restore trainer state if both are available
+    if trainer is not None and schema.trainer_state is not None:
+        ts = schema.trainer_state
+        trainer.current_epoch = ts.get("epoch", 0)
+        trainer.should_stop = ts.get("should_stop", False)
+
+        if "optimizer" in ts:
+            trainer.optimizer.load_state_dict(ts["optimizer"])
+
+        if "scheduler" in ts and trainer.scheduler is not None:
+            trainer.scheduler.load_state_dict(ts["scheduler"])
+
+        if "scaler" in ts and trainer.scaler is not None:
+            trainer.scaler.load_state_dict(ts["scaler"])
+
+        logger.info(
+            "Trainer state restored: epoch=%d  should_stop=%s",
+            trainer.current_epoch,
+            trainer.should_stop,
+        )
+
+    return schema

marn/checkpoint/save.py

@@ -0,0 +1,122 @@
+"""Save a mapping network checkpoint to disk."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from marn.checkpoint.schema import CURRENT_SCHEMA_VERSION, CheckpointSchema
+
+if TYPE_CHECKING:
+    from marn.config.model import MappingConfig
+    from marn.models.mapping_model import MappingModel
+    from marn.trainers.trainer import MappingTrainer
+
+logger = logging.getLogger("marn.checkpoint")
+
+
+def save_checkpoint(
+    path: str | Path,
+    model: MappingModel,
+    *,
+    config: MappingConfig | None = None,
+    trainer: MappingTrainer | None = None,
+    save_trainer_state: bool = True,
+    metadata: dict[str, Any] | None = None,
+) -> None:
+    """Save a compact, versioned checkpoint for the mapping network.
+
+    Only trainable latent vectors and required mapper buffers are written.
+    Generated target weights are excluded — they are ephemeral and can be
+    fully reconstructed at inference time.
+
+    Args:
+        path: Destination file path (e.g. ``"ckpt/epoch10.pt"``).
+            Parent directories are created automatically.
+        model: A ``MappingModel`` instance (may be on any device).
+        config: Optional ``MappingConfig`` saved for reference.  Not used
+            during load, but useful for reproducibility auditing.
+        trainer: Optional ``MappingTrainer`` whose optimizer, scheduler, and
+            epoch state are serialised when ``save_trainer_state`` is ``True``.
+        save_trainer_state: When *True* and a *trainer* is supplied, include
+            optimizer state dict, scheduler state dict (if any), current epoch,
+            and ``should_stop`` flag so training can be resumed exactly.
+        metadata: Free-form dict of user annotations (dataset name, notes …).
+
+    Example::
+
+        save_checkpoint(
+            "checkpoints/epoch10.pt",
+            model=model,
+            config=mapping_config,
+            trainer=trainer,
+        )
+    """
+    import marn  # local import to avoid circular at module level
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # ── Latent state ──────────────────────────────────────────────────────
+    # Only save trainable parameters (latent vectors); mapper buffers are
+    # captured separately so the two sets are clearly delineated.
+    latent_names = {name for name, _ in model.generator.named_latent_vectors()}
+    full_state = model.generator.state_dict()
+
+    latent_state: dict[str, Any] = {k: v for k, v in full_state.items() if k in latent_names}
+    mapper_buffers: dict[str, Any] = {k: v for k, v in full_state.items() if k not in latent_names}
+
+    # ── ParameterSpec fingerprint ─────────────────────────────────────────
+    spec = model.target.parameter_spec
+    spec_names: list[tuple[str, list[int]]] = [(entry.name, list(entry.shape)) for entry in spec]
+
+    # ── Config ───────────────────────────────────────────────────────────
+    config_dict: dict[str, Any] | None = None
+    if config is not None:
+        config_dict = config.model_dump()
+
+    # ── Trainer state ────────────────────────────────────────────────────
+    trainer_state: dict[str, Any] | None = None
+    if trainer is not None and save_trainer_state:
+        trainer_state = {
+            "epoch": trainer.current_epoch,
+            "should_stop": trainer.should_stop,
+            "optimizer": trainer.optimizer.state_dict(),
+        }
+        if trainer.scheduler is not None:
+            trainer_state["scheduler"] = trainer.scheduler.state_dict()
+        if trainer.scaler is not None:
+            trainer_state["scaler"] = trainer.scaler.state_dict()
+
+    schema = CheckpointSchema(
+        schema_version=CURRENT_SCHEMA_VERSION,
+        package_version=marn.__version__,
+        latent_state=latent_state,
+        mapper_buffers=mapper_buffers,
+        parameter_spec_names=spec_names,
+        config=config_dict,
+        trainer_state=trainer_state,
+        metadata=metadata or {},
+    )
+
+    payload = {
+        "schema_version": schema.schema_version,
+        "package_version": schema.package_version,
+        "latent_state": schema.latent_state,
+        "mapper_buffers": schema.mapper_buffers,
+        "parameter_spec_names": schema.parameter_spec_names,
+        "config": schema.config,
+        "trainer_state": schema.trainer_state,
+        "metadata": schema.metadata,
+    }
+
+    torch.save(payload, path)
+    logger.info(
+        "Checkpoint saved to %s  (latents=%d  buffers=%d)",
+        path,
+        len(latent_state),
+        len(mapper_buffers),
+    )

marn/checkpoint/schema.py

@@ -0,0 +1,54 @@
+"""Versioned checkpoint schema for mapping network training state.
+
+The checkpoint stores only what is needed to reconstruct and continue training:
+
+- ``schema_version``: integer bumped on breaking format changes.
+- ``package_version``: the marn package version string at save time.
+- ``latent_state``: ``state_dict`` subset containing only the trainable latent
+  tensors (keys produced by ``model.generator.named_latent_vectors``).
+- ``mapper_buffers``: buffer state required for deterministic reconstruction
+  (e.g. fixed orthogonal projection matrices).
+- ``parameter_spec_names``: ordered parameter names and shapes from the target's
+  ``ParameterSpec``.  Used to validate that the same architecture is loaded.
+- ``config``: optional serialised ``MappingConfig`` dict saved for reference.
+- ``trainer_state``: optional dict containing optimizer/scheduler/epoch state
+  so training can be resumed at the same point.
+- ``metadata``: free-form dict for user annotations (dataset, notes, …).
+
+Generated target weights are **not** included.  Checkpoints are compact because
+only latent vectors and small fixed buffers are saved.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+CURRENT_SCHEMA_VERSION: int = 1
+
+
+@dataclass
+class CheckpointSchema:
+    """All fields written to / read from a checkpoint file.
+
+    Attributes:
+        schema_version: Integer incremented on breaking format changes.
+        package_version: ``marn.__version__`` at save time.
+        latent_state: Ordered ``state_dict`` of trainable latent tensors only.
+        mapper_buffers: Buffer ``state_dict`` required for deterministic
+            reconstruction (projection matrices, running statistics …).
+        parameter_spec_names: Ordered list of ``(name, shape)`` tuples from
+            the target ``ParameterSpec``.  Validated on load.
+        config: Optional serialised ``MappingConfig`` dict (may be ``None``).
+        trainer_state: Optional ``dict`` with optimizer/scheduler/epoch entries.
+        metadata: User-supplied annotations, e.g. dataset or description.
+    """
+
+    schema_version: int
+    package_version: str
+    latent_state: dict[str, Any]
+    mapper_buffers: dict[str, Any]
+    parameter_spec_names: list[tuple[str, list[int]]]
+    config: dict[str, Any] | None = None
+    trainer_state: dict[str, Any] | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)