clip.train_marsclip_mae

train_marsclip_mae

Minimal Stage A MAE training utilities and CLI for MarsCLIP.

WandbLogger dataclass

WandbLogger(run: Any, module: Any, mode: str, project: str, run_name: str | None, log_dir: str)

Small wrapper around an optional W&B run.

resolve_mae_model_config

resolve_mae_model_config(config: dict[str, Any] | None = None, *, preset: str | None = None) -> dict[str, Any]

Resolve checkpoint config values into MAE constructor kwargs.

Source code in src/clip/train_marsclip_mae.py

def resolve_mae_model_config(
    config: dict[str, Any] | None = None,
    *,
    preset: str | None = None,
) -> dict[str, Any]:
    """Resolve checkpoint config values into MAE constructor kwargs."""
    preset_name = (preset or (config or {}).get("model_preset") or "mars_small").strip()
    if preset_name not in MAE_MODEL_PRESETS:
        available = ", ".join(sorted(MAE_MODEL_PRESETS))
        raise ValueError(f"Unknown MAE model preset '{preset_name}'. Available presets: {available}")

    preset_defaults = dict(MAE_MODEL_PRESETS[preset_name])
    source = dict(config or {})

    def _get(key: str, default_key: str | None = None) -> Any:
        value = source.get(key)
        if value is None:
            lookup_key = default_key or key
            return preset_defaults[lookup_key]
        return value

    return {
        "image_size": int(_get("image_size")),
        "patch_size": int(_get("patch_size_px", "patch_size_px") if source.get("patch_size") is None else _get("patch_size")),
        "in_channels": int(_get("in_channels")),
        "encoder_dim": int(_get("encoder_dim")),
        "encoder_depth": int(_get("encoder_depth")),
        "encoder_heads": int(_get("encoder_heads")),
        "decoder_dim": int(_get("decoder_dim")),
        "decoder_depth": int(_get("decoder_depth")),
        "decoder_heads": int(_get("decoder_heads")),
        "min_valid_fraction": float(_get("min_valid_fraction")),
        "normalize_inputs": bool(_get("normalize_inputs")),
        "normalize_targets": bool(_get("normalize_targets")),
        "input_mean": source.get("input_mean", preset_defaults.get("input_mean")),
        "input_std": source.get("input_std", preset_defaults.get("input_std")),
    }

build_mae_model_from_config

build_mae_model_from_config(config: dict[str, Any] | None = None) -> MarsMaskedAutoencoder

Instantiate a Stage A MAE from a saved checkpoint config.

Source code in src/clip/train_marsclip_mae.py

def build_mae_model_from_config(config: dict[str, Any] | None = None) -> MarsMaskedAutoencoder:
    """Instantiate a Stage A MAE from a saved checkpoint config."""
    return MarsMaskedAutoencoder(**resolve_mae_model_config(config))

count_trainable_parameters

count_trainable_parameters(model: Module) -> tuple[int, int]

Return (trainable, total) parameter counts for a model.

Source code in src/clip/train_marsclip_mae.py

def count_trainable_parameters(model: torch.nn.Module) -> tuple[int, int]:
    """Return (trainable, total) parameter counts for a model."""
    total = sum(parameter.numel() for parameter in model.parameters())
    trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
    return trainable, total

compute_valid_pixel_channel_stats

compute_valid_pixel_channel_stats(dataset: Dataset | list[dict[str, Any]], *, max_samples: int | None = None) -> tuple[list[float], list[float]]

Compute per-channel mean/std using only pixels marked valid.

Source code in src/clip/train_marsclip_mae.py

def compute_valid_pixel_channel_stats(
    dataset: Dataset | list[dict[str, Any]],
    *,
    max_samples: int | None = None,
) -> tuple[list[float], list[float]]:
    """Compute per-channel mean/std using only pixels marked valid."""
    if max_samples is not None and max_samples <= 0:
        raise ValueError("max_samples must be positive when provided.")

    sample_count = len(dataset)
    if sample_count == 0:
        raise ValueError("dataset must not be empty.")
    limit = sample_count if max_samples is None else min(sample_count, max_samples)

    channel_sum: torch.Tensor | None = None
    channel_sq_sum: torch.Tensor | None = None
    channel_count: torch.Tensor | None = None

    for idx in range(limit):
        sample = dataset[idx]
        image = torch.as_tensor(sample["image"], dtype=torch.float64)
        valid_mask = torch.as_tensor(sample["valid_mask"], dtype=torch.bool)
        if image.ndim != 3:
            raise ValueError("dataset samples must provide image tensors shaped (C, H, W).")
        if valid_mask.shape != image.shape[-2:]:
            raise ValueError("dataset sample valid_mask must match image spatial dimensions.")
        expanded_valid = valid_mask.unsqueeze(0).expand(image.shape[0], -1, -1)
        masked = torch.where(expanded_valid, image, torch.zeros_like(image))
        valid_counts = expanded_valid.to(torch.float64).sum(dim=(1, 2))

        if channel_sum is None:
            channels = image.shape[0]
            channel_sum = torch.zeros(channels, dtype=torch.float64)
            channel_sq_sum = torch.zeros(channels, dtype=torch.float64)
            channel_count = torch.zeros(channels, dtype=torch.float64)

        channel_sum += masked.sum(dim=(1, 2))
        channel_sq_sum += masked.pow(2).sum(dim=(1, 2))
        channel_count += valid_counts

    assert channel_sum is not None and channel_sq_sum is not None and channel_count is not None
    if not bool((channel_count > 0).all()):
        raise ValueError("All channels must have at least one valid pixel to compute normalization stats.")

    mean = channel_sum / channel_count
    variance = (channel_sq_sum / channel_count) - mean.pow(2)
    std = variance.clamp_min(1e-12).sqrt()
    return mean.to(torch.float32).tolist(), std.to(torch.float32).tolist()

build_optimizer

build_optimizer(model: Module, *, optimizer_name: str = 'adamw', learning_rate: float = 0.0001, weight_decay: float = 0.01) -> Optimizer

Build the requested optimizer for Stage A training.

Source code in src/clip/train_marsclip_mae.py

def build_optimizer(
    model: torch.nn.Module,
    *,
    optimizer_name: str = "adamw",
    learning_rate: float = 1e-4,
    weight_decay: float = 1e-2,
) -> Optimizer:
    """Build the requested optimizer for Stage A training."""
    normalized_name = optimizer_name.lower().strip()
    if normalized_name == "adamw":
        return AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    if normalized_name == "adam":
        return Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    raise ValueError("optimizer_name must be one of: adamw, adam.")

build_scheduler

build_scheduler(optimizer: Optimizer, *, scheduler_name: str = 'none', total_steps: int, warmup_steps: int = 0, min_lr_ratio: float = 0.0) -> LRScheduler | None

Build an optional per-step learning-rate scheduler.

Source code in src/clip/train_marsclip_mae.py

def build_scheduler(
    optimizer: Optimizer,
    *,
    scheduler_name: str = "none",
    total_steps: int,
    warmup_steps: int = 0,
    min_lr_ratio: float = 0.0,
) -> LRScheduler | None:
    """Build an optional per-step learning-rate scheduler."""
    normalized_name = scheduler_name.lower().strip()
    if normalized_name == "none":
        return None
    if total_steps <= 0:
        raise ValueError("total_steps must be positive when a scheduler is enabled.")
    if warmup_steps < 0:
        raise ValueError("warmup_steps must be non-negative.")
    if warmup_steps >= total_steps:
        raise ValueError("warmup_steps must be smaller than total_steps.")
    if not (0.0 <= min_lr_ratio <= 1.0):
        raise ValueError("min_lr_ratio must satisfy 0 <= value <= 1.")

    if normalized_name == "cosine":
        def _lr_lambda(step: int) -> float:
            if warmup_steps > 0 and step < warmup_steps:
                return max(float(step + 1) / float(warmup_steps), 1e-8)
            if total_steps == warmup_steps:
                return min_lr_ratio
            progress = float(step - warmup_steps) / float(max(total_steps - warmup_steps, 1))
            progress = min(max(progress, 0.0), 1.0)
            cosine = 0.5 * (1.0 + torch.cos(torch.tensor(progress * torch.pi)).item())
            return float(min_lr_ratio + (1.0 - min_lr_ratio) * cosine)

        return LambdaLR(optimizer, lr_lambda=_lr_lambda)

    if normalized_name == "onecycle":
        pct_start = float(warmup_steps) / float(total_steps) if warmup_steps > 0 else 0.1
        pct_start = min(max(pct_start, 0.01), 0.99)
        max_lrs = [group["lr"] for group in optimizer.param_groups]
        return OneCycleLR(
            optimizer,
            max_lr=max_lrs,
            total_steps=total_steps,
            pct_start=pct_start,
            anneal_strategy="cos",
        )

    raise ValueError("scheduler_name must be one of: none, cosine, onecycle.")

current_learning_rate

current_learning_rate(optimizer: Optimizer) -> float

Return the learning rate from the first optimizer param group.

Source code in src/clip/train_marsclip_mae.py

def current_learning_rate(optimizer: Optimizer) -> float:
    """Return the learning rate from the first optimizer param group."""
    return float(optimizer.param_groups[0]["lr"])

init_wandb_logger

init_wandb_logger(*, mode: str = 'disabled', project: str = 'marsclip-stagea', run_name: str | None = None, out_dir: Path | str = pathlib.Path('../../tests'), log_dir: Path | str | None = None, config: dict[str, Any] | None = None) -> WandbLogger | None

Initialize an optional W&B run.

The repo defaults to disabled logging so tests and simple runs do not require the dependency. Users can opt into offline mode for local-only dashboards or online when they want to sync to W&B.

Source code in src/clip/train_marsclip_mae.py

def init_wandb_logger(
    *,
    mode: str = "disabled",
    project: str = "marsclip-stagea",
    run_name: str | None = None,
    out_dir: pathlib.Path | str = pathlib.Path("../../tests"),
    log_dir: pathlib.Path | str | None = None,
    config: dict[str, Any] | None = None,
) -> WandbLogger | None:
    """Initialize an optional W&B run.

    The repo defaults to disabled logging so tests and simple runs do not
    require the dependency. Users can opt into `offline` mode for local-only
    dashboards or `online` when they want to sync to W&B.
    """
    normalized_mode = mode.lower().strip()
    if normalized_mode == "disabled":
        return None
    if normalized_mode not in {"offline", "online"}:
        raise ValueError("W&B mode must be one of: disabled, offline, online.")
    try:
        wandb = importlib.import_module("wandb")
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "wandb is not installed. Install it before using --wandb-mode offline or online."
        ) from exc

    resolved_log_dir = pathlib.Path(log_dir) if log_dir is not None else pathlib.Path(out_dir) / "wandb"
    resolved_log_dir.mkdir(parents=True, exist_ok=True)
    run = wandb.init(
        project=project,
        name=run_name,
        mode=normalized_mode,
        dir=str(resolved_log_dir),
        config=dict(config or {}),
        reinit="finish_previous",
    )
    return WandbLogger(
        run=run,
        module=wandb,
        mode=normalized_mode,
        project=project,
        run_name=run_name,
        log_dir=str(resolved_log_dir),
    )

build_mae_dataloader

build_mae_dataloader(dataset: Dataset | list[dict[str, Any]], *, batch_size: int = 4, shuffle: bool = True, generator: Generator | None = None, num_workers: int = 0, pin_memory: bool = False, prefetch_factor: int | None = None, persistent_workers: bool = False) -> DataLoader

Build a DataLoader that emits Stage A MAE-ready patch batches.

Source code in src/clip/train_marsclip_mae.py

def build_mae_dataloader(
    dataset: Dataset | list[dict[str, Any]],
    *,
    batch_size: int = 4,
    shuffle: bool = True,
    generator: torch.Generator | None = None,
    num_workers: int = 0,
    pin_memory: bool = False,
    prefetch_factor: int | None = None,
    persistent_workers: bool = False,
) -> DataLoader:
    """Build a DataLoader that emits Stage A MAE-ready patch batches."""
    kwargs: dict[str, Any] = {
        "dataset": dataset,
        "batch_size": batch_size,
        "shuffle": shuffle,
        "generator": generator,
        "num_workers": num_workers,
        "pin_memory": pin_memory,
        "collate_fn": collate_patch_samples_for_mae,
    }
    if num_workers > 0:
        kwargs["persistent_workers"] = persistent_workers
        if prefetch_factor is not None:
            kwargs["prefetch_factor"] = prefetch_factor
    return DataLoader(**kwargs)

resolve_map_location

resolve_map_location(map_location: str | device = 'cpu') -> str | torch.device

Normalize torch.load map_location values, including the repo's 'auto' alias.

Source code in src/clip/train_marsclip_mae.py

def resolve_map_location(map_location: str | torch.device = "cpu") -> str | torch.device:
    """Normalize torch.load map_location values, including the repo's 'auto' alias."""
    if isinstance(map_location, str) and map_location.lower() == "auto":
        return "cuda" if torch.cuda.is_available() else "cpu"
    return map_location

train_mae_batch

train_mae_batch(model: MarsMaskedAutoencoder, batch: dict[str, Any], optimizer: Optimizer, *, mask_ratio: float = 0.75, device: str | device | None = None, step: int | None = None) -> tuple[dict[str, float], MarsMAEOutput]

Run one MAE optimization step on a collated batch.

Source code in src/clip/train_marsclip_mae.py

def train_mae_batch(
    model: MarsMaskedAutoencoder,
    batch: dict[str, Any],
    optimizer: Optimizer,
    *,
    mask_ratio: float = 0.75,
    device: str | torch.device | None = None,
    step: int | None = None,
) -> tuple[dict[str, float], MarsMAEOutput]:
    """Run one MAE optimization step on a collated batch."""
    resolved_device = _resolve_device(device, model)
    model.to(resolved_device)
    model.train()

    batch = _move_batch_to_device(batch, resolved_device)
    optimizer.zero_grad(set_to_none=True)
    output = model.forward_patch_batch(batch, mask_ratio=mask_ratio)
    loss = output.loss
    if not torch.isfinite(loss):
        raise ValueError("Encountered non-finite MAE loss.")
    loss.backward()
    optimizer.step()

    entry = {
        "loss": float(loss.detach().cpu()),
        "visible_patch_fraction": float(output.visible_mask.float().mean().detach().cpu()),
        "masked_valid_patch_fraction": float(output.masked_valid_mask.float().mean().detach().cpu()),
        "loss_pixel_fraction": float(output.loss_mask.float().mean().detach().cpu()),
    }
    if step is not None:
        entry["step"] = float(step)
    return entry, _output_to_cpu(output)

train_mae_steps

train_mae_steps(model: MarsMaskedAutoencoder, dataloader: DataLoader, optimizer: Optimizer, *, num_steps: int, mask_ratio: float = 0.75, device: str | device | None = None, start_step: int = 0, step_callback: Callable[[int, dict[str, float]], None] | None = None) -> list[dict[str, float]]

Run a fixed number of MAE optimization steps, cycling the dataloader if needed.

Source code in src/clip/train_marsclip_mae.py

def train_mae_steps(
    model: MarsMaskedAutoencoder,
    dataloader: DataLoader,
    optimizer: Optimizer,
    *,
    num_steps: int,
    mask_ratio: float = 0.75,
    device: str | torch.device | None = None,
    start_step: int = 0,
    step_callback: Callable[[int, dict[str, float]], None] | None = None,
) -> list[dict[str, float]]:
    """Run a fixed number of MAE optimization steps, cycling the dataloader if needed."""
    if num_steps < 0:
        raise ValueError("num_steps must be non-negative.")

    resolved_device = _resolve_device(device, model)
    model.to(resolved_device)
    model.train()

    history: list[dict[str, float]] = []
    iterator: Iterable[Any] | Any = iter(dataloader)
    for step_idx in range(num_steps):
        try:
            batch = next(iterator)
        except StopIteration:
            iterator = iter(dataloader)
            batch = next(iterator)

        step = start_step + step_idx + 1
        entry, _ = train_mae_batch(
            model,
            batch,
            optimizer,
            mask_ratio=mask_ratio,
            device=resolved_device,
            step=step,
        )
        history.append(entry)
        if step_callback is not None:
            step_callback(step, entry)
    return history

evaluate_mae_dataloader

evaluate_mae_dataloader(model: MarsMaskedAutoencoder, dataloader: DataLoader, *, mask_ratio: float = 0.75, device: str | device | None = None, max_batches: int | None = None, batch_callback: Callable[[int, int | None], None] | None = None) -> dict[str, float]

Evaluate MAE reconstruction loss on a validation/test dataloader.

Source code in src/clip/train_marsclip_mae.py

def evaluate_mae_dataloader(
    model: MarsMaskedAutoencoder,
    dataloader: DataLoader,
    *,
    mask_ratio: float = 0.75,
    device: str | torch.device | None = None,
    max_batches: int | None = None,
    batch_callback: Callable[[int, int | None], None] | None = None,
) -> dict[str, float]:
    """Evaluate MAE reconstruction loss on a validation/test dataloader."""
    if max_batches is not None and max_batches <= 0:
        raise ValueError("max_batches must be positive when provided.")

    resolved_device = _resolve_device(device, model)
    model.to(resolved_device)
    was_training = model.training
    model.eval()

    total_loss = 0.0
    total_visible = 0.0
    total_masked = 0.0
    total_loss_pixels = 0.0
    num_batches = 0
    total_batches: int | None
    try:
        total_batches = len(dataloader)
    except TypeError:
        total_batches = None
    if max_batches is not None:
        total_batches = min(total_batches, max_batches) if total_batches is not None else max_batches

    with torch.no_grad():
        for batch_idx, batch in enumerate(dataloader):
            if max_batches is not None and batch_idx >= max_batches:
                break
            if batch_callback is not None:
                batch_callback(batch_idx + 1, total_batches)
            batch = _move_batch_to_device(batch, resolved_device)
            output = model.forward_patch_batch(batch, mask_ratio=mask_ratio)
            total_loss += float(output.loss.detach().cpu())
            total_visible += float(output.visible_mask.float().mean().detach().cpu())
            total_masked += float(output.masked_valid_mask.float().mean().detach().cpu())
            total_loss_pixels += float(output.loss_mask.float().mean().detach().cpu())
            num_batches += 1

    if was_training:
        model.train()
    if num_batches == 0:
        raise ValueError("Validation/test dataloader produced no batches.")

    return {
        "loss": total_loss / float(num_batches),
        "visible_patch_fraction": total_visible / float(num_batches),
        "masked_valid_patch_fraction": total_masked / float(num_batches),
        "loss_pixel_fraction": total_loss_pixels / float(num_batches),
        "num_batches": float(num_batches),
    }

save_mae_checkpoint

save_mae_checkpoint(path: Path | str, model: MarsMaskedAutoencoder, optimizer: Optimizer | None, scheduler: LRScheduler | None = None, *, step: int, history: list[dict[str, float]], config: dict[str, Any] | None = None) -> pathlib.Path

Save a minimal MAE training checkpoint.

Source code in src/clip/train_marsclip_mae.py

def save_mae_checkpoint(
    path: pathlib.Path | str,
    model: MarsMaskedAutoencoder,
    optimizer: Optimizer | None,
    scheduler: LRScheduler | None = None,
    *,
    step: int,
    history: list[dict[str, float]],
    config: dict[str, Any] | None = None,
) -> pathlib.Path:
    """Save a minimal MAE training checkpoint."""
    out = pathlib.Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "step": int(step),
            "history": history,
            "config": config or {},
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict() if optimizer is not None else None,
            "scheduler_state": scheduler.state_dict() if scheduler is not None else None,
        },
        out,
    )
    return out

load_mae_checkpoint

load_mae_checkpoint(path: Path | str, model: MarsMaskedAutoencoder, optimizer: Optimizer | None = None, scheduler: LRScheduler | None = None, *, map_location: str | device = 'cpu') -> dict[str, Any]

Load a MAE checkpoint into a model and optional optimizer.

Source code in src/clip/train_marsclip_mae.py

def load_mae_checkpoint(
    path: pathlib.Path | str,
    model: MarsMaskedAutoencoder,
    optimizer: Optimizer | None = None,
    scheduler: LRScheduler | None = None,
    *,
    map_location: str | torch.device = "cpu",
) -> dict[str, Any]:
    """Load a MAE checkpoint into a model and optional optimizer."""
    checkpoint = torch.load(path, map_location=resolve_map_location(map_location))
    model.load_state_dict(checkpoint["model_state"])
    if optimizer is not None and checkpoint.get("optimizer_state") is not None:
        optimizer.load_state_dict(checkpoint["optimizer_state"])
    if scheduler is not None and checkpoint.get("scheduler_state") is not None:
        scheduler.load_state_dict(checkpoint["scheduler_state"])
    return {
        "step": int(checkpoint.get("step", 0)),
        "history": list(checkpoint.get("history", [])),
        "config": dict(checkpoint.get("config", {})),
    }

save_training_history

save_training_history(history: list[dict[str, float]], out_path: Path | str) -> pathlib.Path

Persist MAE loss history as JSON.

Source code in src/clip/train_marsclip_mae.py

def save_training_history(
    history: list[dict[str, float]],
    out_path: pathlib.Path | str,
) -> pathlib.Path:
    """Persist MAE loss history as JSON."""
    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(history, indent=2))
    return out

select_best_history_entry

select_best_history_entry(history: list[dict[str, float]]) -> dict[str, float]

Return the minimum-loss training record.

Source code in src/clip/train_marsclip_mae.py

def select_best_history_entry(history: list[dict[str, float]]) -> dict[str, float]:
    """Return the minimum-loss training record."""
    if not history:
        raise ValueError("history must not be empty.")
    return min(history, key=lambda item: float(item["loss"]))

save_training_summary

save_training_summary(summary: dict[str, Any], out_path: Path | str) -> pathlib.Path

Persist the MAE training summary as JSON.

Source code in src/clip/train_marsclip_mae.py

def save_training_summary(
    summary: dict[str, Any],
    out_path: pathlib.Path | str,
) -> pathlib.Path:
    """Persist the MAE training summary as JSON."""
    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(summary, indent=2))
    return out

save_training_progress

save_training_progress(progress: dict[str, Any], out_path: Path | str) -> pathlib.Path

Persist an incremental MAE training progress snapshot as JSON.

Source code in src/clip/train_marsclip_mae.py

def save_training_progress(
    progress: dict[str, Any],
    out_path: pathlib.Path | str,
) -> pathlib.Path:
    """Persist an incremental MAE training progress snapshot as JSON."""
    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    payload = dict(progress)
    payload["updated_at"] = _progress_timestamp()
    out.write_text(json.dumps(payload, indent=2))
    return out

load_normalization_stats

load_normalization_stats(path: Path | str) -> tuple[list[float], list[float], dict[str, Any]]

Load cached normalization statistics from JSON.

Source code in src/clip/train_marsclip_mae.py

def load_normalization_stats(path: pathlib.Path | str) -> tuple[list[float], list[float], dict[str, Any]]:
    """Load cached normalization statistics from JSON."""
    source = pathlib.Path(path)
    payload = json.loads(source.read_text())
    mean = [float(value) for value in payload["mean"]]
    std = [float(value) for value in payload["std"]]
    return mean, std, payload

save_normalization_stats

save_normalization_stats(path: Path | str, *, mean: list[float], std: list[float], metadata: dict[str, Any] | None = None) -> pathlib.Path

Persist normalization statistics for reuse across comparable runs.

Source code in src/clip/train_marsclip_mae.py

def save_normalization_stats(
    path: pathlib.Path | str,
    *,
    mean: list[float],
    std: list[float],
    metadata: dict[str, Any] | None = None,
) -> pathlib.Path:
    """Persist normalization statistics for reuse across comparable runs."""
    out = pathlib.Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "mean": [float(value) for value in mean],
        "std": [float(value) for value in std],
    }
    if metadata:
        payload.update(dict(metadata))
    out.write_text(json.dumps(payload, indent=2))
    return out

save_loss_curve

save_loss_curve(history: list[dict[str, float]], out_path: Path | str) -> pathlib.Path

Save a simple MAE loss curve PNG using a log-scaled y-axis.

Source code in src/clip/train_marsclip_mae.py

def save_loss_curve(
    history: list[dict[str, float]],
    out_path: pathlib.Path | str,
) -> pathlib.Path:
    """Save a simple MAE loss curve PNG using a log-scaled y-axis."""
    if not history:
        raise ValueError("history must not be empty.")

    steps = [entry["step"] for entry in history]
    losses = [entry["loss"] for entry in history]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(steps, losses, color="#d95f02", linewidth=2)
    ax.set_xlabel("Step")
    ax.set_ylabel("MAE loss")
    ax.set_yscale("log")
    ax.set_title("MarsCLIP Stage A training loss")
    ax.grid(alpha=0.25)
    fig.tight_layout()

    out = pathlib.Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=160)
    plt.close(fig)
    return out

generate_mae_preview

generate_mae_preview(model: MarsMaskedAutoencoder, samples: list[dict[str, Any]], out_path: Path | str, *, mask_ratio: float, device: str | device | None = None) -> pathlib.Path

Generate a masked/reconstructed preview from a small sample set.

Source code in src/clip/train_marsclip_mae.py

def generate_mae_preview(
    model: MarsMaskedAutoencoder,
    samples: list[dict[str, Any]],
    out_path: pathlib.Path | str,
    *,
    mask_ratio: float,
    device: str | torch.device | None = None,
) -> pathlib.Path:
    """Generate a masked/reconstructed preview from a small sample set."""
    if not samples:
        raise ValueError("samples must not be empty.")

    resolved_device = _resolve_device(device, model)
    was_training = model.training
    batch = collate_patch_samples_for_mae(samples)
    batch = _move_batch_to_device(batch, resolved_device)
    model.to(resolved_device)
    model.eval()
    with torch.no_grad():
        output = model.forward_patch_batch(
            batch,
            mask_ratio=mask_ratio,
            generator=torch.Generator().manual_seed(0),
        )
    cpu_output = _output_to_cpu(output)
    if was_training:
        model.train()

    return save_mae_reconstruction_preview(
        samples,
        cpu_output,
        out_path,
        patch_size=model.patch_size,
        max_items=len(samples),
    )

run_mae_training

run_mae_training(model: MarsMaskedAutoencoder, dataloader: DataLoader, optimizer: Optimizer, scheduler: LRScheduler | None = None, *, out_dir: Path | str, num_steps: int, mask_ratio: float, preview_samples: list[dict[str, Any]] | None = None, checkpoint_every: int = 0, preview_every: int = 0, device: str | device | None = None, resume_state: dict[str, Any] | None = None, config: dict[str, Any] | None = None, wandb_logger: WandbLogger | None = None, val_dataloader: DataLoader | None = None, val_every: int = 0, val_max_batches: int | None = None) -> dict[str, Any]

Run a small Stage A training session with periodic artifacts.

Source code in src/clip/train_marsclip_mae.py

def run_mae_training(
    model: MarsMaskedAutoencoder,
    dataloader: DataLoader,
    optimizer: Optimizer,
    scheduler: LRScheduler | None = None,
    *,
    out_dir: pathlib.Path | str,
    num_steps: int,
    mask_ratio: float,
    preview_samples: list[dict[str, Any]] | None = None,
    checkpoint_every: int = 0,
    preview_every: int = 0,
    device: str | torch.device | None = None,
    resume_state: dict[str, Any] | None = None,
    config: dict[str, Any] | None = None,
    wandb_logger: WandbLogger | None = None,
    val_dataloader: DataLoader | None = None,
    val_every: int = 0,
    val_max_batches: int | None = None,
) -> dict[str, Any]:
    """Run a small Stage A training session with periodic artifacts."""
    out_path = pathlib.Path(out_dir)
    out_path.mkdir(parents=True, exist_ok=True)

    history = list((resume_state or {}).get("history", []))
    start_step = int((resume_state or {}).get("step", 0))
    checkpoint_paths: list[str] = []
    preview_paths: list[str] = []
    progress_path = out_path / "progress.json"

    best_entry = select_best_history_entry(history) if history and val_dataloader is None else None
    best_loss = float(best_entry["loss"]) if best_entry is not None else float("inf")
    best_step = int(best_entry["step"]) if best_entry is not None else start_step
    best_metric_name = "val_loss" if val_dataloader is not None else "train_loss"
    best_checkpoint_path = out_path / "best_checkpoint.pt"
    val_history: list[dict[str, float]] = []

    save_training_progress(
        {
            "status": "running",
            "phase": "starting",
            "start_step": start_step,
            "current_step": start_step,
            "target_step": start_step + num_steps,
            "best_loss": None,
            "best_step": start_step,
            "latest_loss": None,
            "latest_lr": current_learning_rate(optimizer),
            "val_best_loss": None,
            "out_dir": str(out_path),
        },
        progress_path,
    )

    start_preview = None
    if preview_samples:
        start_preview = generate_mae_preview(
            model,
            preview_samples,
            out_path / f"preview_step_{start_step:06d}.png",
            mask_ratio=mask_ratio,
            device=device,
        )
        preview_paths.append(str(start_preview))
        if wandb_logger is not None:
            wandb_logger.log_image(
                "train/preview_start",
                start_preview,
                step=start_step,
                caption=f"Stage A preview at step {start_step}",
            )
        save_training_progress(
            {
                "status": "running",
                "phase": "preview_start",
                "start_step": start_step,
                "current_step": start_step,
                "target_step": start_step + num_steps,
                "best_loss": best_loss if best_loss != float("inf") else None,
                "best_step": best_step,
                "latest_loss": None,
                "latest_lr": current_learning_rate(optimizer),
                "val_best_loss": None,
                "out_dir": str(out_path),
                "latest_preview": str(start_preview),
            },
            progress_path,
        )

    iterator = iter(dataloader)
    resolved_device = _resolve_device(device, model)
    for offset in range(num_steps):
        try:
            batch = next(iterator)
        except StopIteration:
            iterator = iter(dataloader)
            batch = next(iterator)

        step = start_step + offset + 1
        entry, _ = train_mae_batch(
            model,
            batch,
            optimizer,
            mask_ratio=mask_ratio,
            device=resolved_device,
            step=step,
        )
        if scheduler is not None:
            scheduler.step()
        entry["lr"] = current_learning_rate(optimizer)
        history.append(entry)
        if wandb_logger is not None:
            wandb_logger.log_metrics(
                {
                    "train/loss": float(entry["loss"]),
                    "train/lr": float(entry["lr"]),
                    "train/visible_patch_fraction": float(entry["visible_patch_fraction"]),
                    "train/masked_valid_patch_fraction": float(entry["masked_valid_patch_fraction"]),
                    "train/loss_pixel_fraction": float(entry["loss_pixel_fraction"]),
                },
                step=step,
            )
        current_val_best = min((float(item["loss"]) for item in val_history), default=None)
        save_training_progress(
            {
                "status": "running",
                "phase": "training",
                "start_step": start_step,
                "current_step": step,
                "target_step": start_step + num_steps,
                "best_loss": best_loss if best_loss != float("inf") else None,
                "best_step": best_step,
                "latest_loss": float(entry["loss"]),
                "latest_lr": float(entry["lr"]),
                "val_best_loss": current_val_best,
                "out_dir": str(out_path),
            },
            progress_path,
        )

        if val_dataloader is None and float(entry["loss"]) <= best_loss:
            best_loss = float(entry["loss"])
            best_step = step
            save_mae_checkpoint(
                best_checkpoint_path,
                model,
                optimizer,
                scheduler,
                step=step,
                history=history,
                config=config,
            )
            if wandb_logger is not None:
                wandb_logger.log_metrics(
                    {
                        "train/best_loss": float(best_loss),
                        "train/best_step": float(best_step),
                        "train/best_lr": float(entry["lr"]),
                    },
                    step=step,
                )

        should_run_val = val_dataloader is not None and (
            (val_every > 0 and step % val_every == 0)
            or (val_every <= 0 and checkpoint_every > 0 and step % checkpoint_every == 0)
        )
        if should_run_val:
            try:
                expected_val_batches = len(val_dataloader)
            except TypeError:
                expected_val_batches = None
            if val_max_batches is not None:
                expected_val_batches = (
                    min(expected_val_batches, val_max_batches)
                    if expected_val_batches is not None
                    else val_max_batches
                )
            save_training_progress(
                {
                    "status": "running",
                    "phase": "validating",
                    "start_step": start_step,
                    "current_step": step,
                    "target_step": start_step + num_steps,
                    "best_loss": best_loss if best_loss != float("inf") else None,
                    "best_step": best_step,
                    "latest_loss": float(entry["loss"]),
                    "latest_lr": float(entry["lr"]),
                    "val_progress_batches_completed": 0,
                    "val_progress_batches_total": expected_val_batches,
                    "val_best_loss": min((float(item["loss"]) for item in val_history), default=None),
                    "out_dir": str(out_path),
                },
                progress_path,
            )

            def _on_val_batch(batch_number: int, total_batches: int | None) -> None:
                if batch_number == 1 or batch_number % 5 == 0 or batch_number == total_batches:
                    save_training_progress(
                        {
                            "status": "running",
                            "phase": "validating",
                            "start_step": start_step,
                            "current_step": step,
                            "target_step": start_step + num_steps,
                            "best_loss": best_loss if best_loss != float("inf") else None,
                            "best_step": best_step,
                            "latest_loss": float(entry["loss"]),
                            "latest_lr": float(entry["lr"]),
                            "val_progress_batches_completed": batch_number,
                            "val_progress_batches_total": total_batches,
                            "val_best_loss": min((float(item["loss"]) for item in val_history), default=None),
                            "out_dir": str(out_path),
                        },
                        progress_path,
                    )

            val_entry = evaluate_mae_dataloader(
                model,
                val_dataloader,
                mask_ratio=mask_ratio,
                device=resolved_device,
                max_batches=val_max_batches,
                batch_callback=_on_val_batch,
            )
            val_entry["step"] = float(step)
            val_history.append(val_entry)
            if wandb_logger is not None:
                wandb_logger.log_metrics(
                    {
                        "val/loss": float(val_entry["loss"]),
                        "val/visible_patch_fraction": float(val_entry["visible_patch_fraction"]),
                        "val/masked_valid_patch_fraction": float(val_entry["masked_valid_patch_fraction"]),
                        "val/loss_pixel_fraction": float(val_entry["loss_pixel_fraction"]),
                    },
                    step=step,
                )
            save_training_progress(
                {
                    "status": "running",
                    "phase": "validating",
                    "start_step": start_step,
                    "current_step": step,
                    "target_step": start_step + num_steps,
                    "best_loss": best_loss if best_loss != float("inf") else None,
                    "best_step": best_step,
                    "latest_loss": float(entry["loss"]),
                    "latest_lr": float(entry["lr"]),
                    "val_latest_loss": float(val_entry["loss"]),
                    "val_best_loss": min(float(item["loss"]) for item in val_history),
                    "out_dir": str(out_path),
                },
                progress_path,
            )
            if float(val_entry["loss"]) <= best_loss:
                best_loss = float(val_entry["loss"])
                best_step = step
                save_mae_checkpoint(
                    best_checkpoint_path,
                    model,
                    optimizer,
                    scheduler,
                    step=step,
                    history=history,
                    config=config,
                )
                if wandb_logger is not None:
                    wandb_logger.log_metrics(
                        {
                            "val/best_loss": float(best_loss),
                            "val/best_step": float(best_step),
                        },
                        step=step,
                    )

        if checkpoint_every > 0 and step % checkpoint_every == 0:
            ckpt = save_mae_checkpoint(
                out_path / f"checkpoint_step_{step:06d}.pt",
                model,
                optimizer,
                scheduler,
                step=step,
                history=history,
                config=config,
            )
            checkpoint_paths.append(str(ckpt))

        if preview_samples and preview_every > 0 and step % preview_every == 0:
            preview = generate_mae_preview(
                model,
                preview_samples,
                out_path / f"preview_step_{step:06d}.png",
                mask_ratio=mask_ratio,
                device=resolved_device,
            )
            preview_paths.append(str(preview))
            if wandb_logger is not None:
                wandb_logger.log_image(
                    "train/preview_periodic",
                    preview,
                    step=step,
                    caption=f"Stage A preview at step {step}",
                )

    final_step = start_step + num_steps
    final_preview = None
    if preview_samples:
        final_preview = generate_mae_preview(
            model,
            preview_samples,
            out_path / "preview_final.png",
            mask_ratio=mask_ratio,
            device=resolved_device,
        )
        preview_paths.append(str(final_preview))
        if wandb_logger is not None:
            wandb_logger.log_image(
                "train/preview_final",
                final_preview,
                step=final_step,
                caption=f"Stage A preview at step {final_step}",
            )

    checkpoint_path = save_mae_checkpoint(
        out_path / "checkpoint.pt",
        model,
        optimizer,
        scheduler,
        step=final_step,
        history=history,
        config=config,
    )
    final_val_entry = None
    if val_dataloader is not None and (not val_history or int(val_history[-1]["step"]) != final_step):
        final_val_entry = evaluate_mae_dataloader(
            model,
            val_dataloader,
            mask_ratio=mask_ratio,
            device=resolved_device,
            max_batches=val_max_batches,
        )
        final_val_entry["step"] = float(final_step)
        val_history.append(final_val_entry)
        if wandb_logger is not None:
            wandb_logger.log_metrics(
                {
                    "val/loss": float(final_val_entry["loss"]),
                    "val/visible_patch_fraction": float(final_val_entry["visible_patch_fraction"]),
                    "val/masked_valid_patch_fraction": float(final_val_entry["masked_valid_patch_fraction"]),
                    "val/loss_pixel_fraction": float(final_val_entry["loss_pixel_fraction"]),
                },
                step=final_step,
            )
        if float(final_val_entry["loss"]) <= best_loss:
            best_loss = float(final_val_entry["loss"])
            best_step = final_step
            save_mae_checkpoint(
                best_checkpoint_path,
                model,
                optimizer,
                scheduler,
                step=final_step,
                history=history,
                config=config,
            )

    history_path = save_training_history(history, out_path / "history.json")
    curve_path = save_loss_curve(history, out_path / "loss_curve.png")
    val_history_path = save_training_history(val_history, out_path / "val_history.json") if val_history else None

    best_metric_label = "training loss" if best_metric_name == "train_loss" else "val loss"

    summary = {
        "start_step": start_step,
        "num_new_steps": num_steps,
        "num_steps_total": len(history),
        "model_preset": (config or {}).get("model_preset"),
        "trainable_parameters": int((config or {}).get("trainable_parameters"))
        if (config or {}).get("trainable_parameters") is not None
        else None,
        "total_parameters": int((config or {}).get("total_parameters"))
        if (config or {}).get("total_parameters") is not None
        else None,
        "batch_size": int((config or {}).get("batch_size")) if (config or {}).get("batch_size") is not None else None,
        "color_only": bool((config or {}).get("color_only"))
        if (config or {}).get("color_only") is not None
        else None,
        "optimizer": (config or {}).get("optimizer"),
        "scheduler": (config or {}).get("scheduler"),
        "normalization_mode": (config or {}).get("normalization_mode"),
        "normalization_max_samples": int((config or {}).get("normalization_max_samples"))
        if (config or {}).get("normalization_max_samples") is not None
        else None,
        "normalization_stats_path": str((config or {}).get("normalization_stats_path"))
        if (config or {}).get("normalization_stats_path") is not None
        else None,
        "patch_records_path": str((config or {}).get("patch_records_path"))
        if (config or {}).get("patch_records_path") is not None
        else None,
        "normalize_inputs": bool((config or {}).get("normalize_inputs"))
        if (config or {}).get("normalize_inputs") is not None
        else None,
        "normalize_targets": bool((config or {}).get("normalize_targets"))
        if (config or {}).get("normalize_targets") is not None
        else None,
        "input_mean": list((config or {}).get("input_mean"))
        if (config or {}).get("input_mean") is not None
        else None,
        "input_std": list((config or {}).get("input_std"))
        if (config or {}).get("input_std") is not None
        else None,
        "scheduler_total_steps": int((config or {}).get("scheduler_total_steps"))
        if (config or {}).get("scheduler_total_steps") is not None
        else None,
        "warmup_steps": int((config or {}).get("warmup_steps")) if (config or {}).get("warmup_steps") is not None else None,
        "min_lr_ratio": float((config or {}).get("min_lr_ratio"))
        if (config or {}).get("min_lr_ratio") is not None
        else None,
        "split_mode": (config or {}).get("split_mode"),
        "split_manifest": str((config or {}).get("split_manifest"))
        if (config or {}).get("split_manifest") is not None
        else None,
        "train_count": int((config or {}).get("train_count")) if (config or {}).get("train_count") is not None else None,
        "val_count": int((config or {}).get("val_count")) if (config or {}).get("val_count") is not None else None,
        "test_count": int((config or {}).get("test_count")) if (config or {}).get("test_count") is not None else None,
        "split_summary_path": str((config or {}).get("split_summary_path"))
        if (config or {}).get("split_summary_path") is not None
        else None,
        "num_workers": int((config or {}).get("num_workers")) if (config or {}).get("num_workers") is not None else None,
        "pin_memory": bool((config or {}).get("pin_memory")) if (config or {}).get("pin_memory") is not None else None,
        "prefetch_factor": int((config or {}).get("prefetch_factor")) if (config or {}).get("prefetch_factor") is not None else None,
        "persistent_workers": bool((config or {}).get("persistent_workers")) if (config or {}).get("persistent_workers") is not None else None,
        "initial_loss": history[0]["loss"] if history else None,
        "final_loss": history[-1]["loss"] if history else None,
        "initial_lr": history[0]["lr"] if history and "lr" in history[0] else current_learning_rate(optimizer),
        "final_lr": history[-1]["lr"] if history and "lr" in history[-1] else current_learning_rate(optimizer),
        "best_loss": best_loss if history else None,
        "best_metric_name": best_metric_name,
        "best_step": best_step if history else None,
        "val_initial_loss": val_history[0]["loss"] if val_history else None,
        "val_final_loss": val_history[-1]["loss"] if val_history else None,
        "val_best_loss": min(float(entry["loss"]) for entry in val_history) if val_history else None,
        "val_history_path": str(val_history_path) if val_history_path is not None else None,
        "val_every": int(val_every) if val_every is not None else None,
        "val_max_batches": int(val_max_batches) if val_max_batches is not None else None,
        "resume_step": start_step if resume_state is not None else None,
        "resume_checkpoint": str((resume_state or {}).get("checkpoint_path"))
        if resume_state is not None and (resume_state or {}).get("checkpoint_path") is not None
        else None,
        "start_preview": str(start_preview) if start_preview is not None else None,
        "final_preview": str(final_preview) if final_preview is not None else None,
        "preview_paths": preview_paths,
        "checkpoint": str(checkpoint_path),
        "checkpoint_paths": checkpoint_paths,
        "best_checkpoint": str(best_checkpoint_path) if history else None,
        "best_checkpoint_rule": f"minimum observed {best_metric_label}",
        "history_path": str(history_path),
        "loss_curve": str(curve_path),
        "progress_path": str(progress_path),
        "wandb_mode": wandb_logger.mode if wandb_logger is not None else "disabled",
        "wandb_project": wandb_logger.project if wandb_logger is not None else None,
        "wandb_run_name": wandb_logger.run_name if wandb_logger is not None else None,
        "wandb_run_id": wandb_logger.run_id if wandb_logger is not None else None,
        "wandb_log_dir": wandb_logger.log_dir if wandb_logger is not None else None,
        "wandb_run_dir": wandb_logger.run_dir if wandb_logger is not None else None,
    }
    if wandb_logger is not None:
        wandb_logger.log_image(
            "train/loss_curve",
            curve_path,
            step=final_step,
            caption="MarsCLIP Stage A training loss curve",
        )
        wandb_logger.finish(summary)
    summary_path = save_training_summary(summary, out_path / "summary.json")
    save_training_progress(
        {
            "status": "completed",
            "phase": "finished",
            "start_step": start_step,
            "current_step": final_step,
            "target_step": final_step,
            "best_loss": best_loss if history else None,
            "best_step": best_step if history else None,
            "latest_loss": history[-1]["loss"] if history else None,
            "latest_lr": history[-1]["lr"] if history and "lr" in history[-1] else current_learning_rate(optimizer),
            "val_best_loss": min((float(item["loss"]) for item in val_history), default=None),
            "out_dir": str(out_path),
            "summary_path": str(summary_path),
            "best_checkpoint": str(best_checkpoint_path) if history else None,
            "final_preview": str(final_preview) if final_preview is not None else None,
        },
        progress_path,
    )
    summary["summary_path"] = str(summary_path)
    return summary