code/dispnet.v2.3.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
dispnet.v2.3.py

残差保持 1D CNN 版本：
- 无 VAE / KL / NLL
- 不输入震中距、period 等额外条件
- 只输入 waveform
- 模型内部保存训练集平均频散 reference_disp
- 网络预测 delta_v = v - reference_disp，再还原得到 disp_mu
- loss 额外加入 batch 内 pairwise 差异保持和方差保持，抑制回归均值化

设计目标：
v2.2 的点误差较低，但反演结果过平滑。v2.3 不再只奖励逐点接近标签，
而是显式要求预测频散保留样本间的横向差异和速度分布宽度。
"""

import os
import random
import time
from dataclasses import dataclass
from typing import Dict, Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader


def set_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def exists(x):
    return x is not None


def default_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def _gn_groups(channels: int, max_groups: int = 8) -> int:
    for g in range(min(max_groups, channels), 0, -1):
        if channels % g == 0:
            return g
    return 1


class ConvGNAct(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 7, stride: int = 1, dilation: int = 1):
        super().__init__()
        padding = dilation * (kernel_size - 1) // 2
        self.net = nn.Sequential(
            nn.Conv1d(
                in_ch,
                out_ch,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                dilation=dilation,
                bias=False,
            ),
            nn.GroupNorm(_gn_groups(out_ch), out_ch),
            nn.SiLU(inplace=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class ResidualBlock(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, stride: int = 1, kernel_size: int = 7, dilation: int = 1):
        super().__init__()
        self.conv1 = ConvGNAct(in_ch, out_ch, kernel_size=kernel_size, stride=stride, dilation=dilation)
        self.conv2 = nn.Sequential(
            nn.Conv1d(
                out_ch,
                out_ch,
                kernel_size=kernel_size,
                stride=1,
                padding=dilation * (kernel_size - 1) // 2,
                dilation=dilation,
                bias=False,
            ),
            nn.GroupNorm(_gn_groups(out_ch), out_ch),
        )
        if stride != 1 or in_ch != out_ch:
            self.shortcut = nn.Sequential(
                nn.Conv1d(in_ch, out_ch, kernel_size=1, stride=stride, bias=False),
                nn.GroupNorm(_gn_groups(out_ch), out_ch),
            )
        else:
            self.shortcut = nn.Identity()
        self.act = nn.SiLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.act(self.conv2(self.conv1(x)) + self.shortcut(x))


class MultiScaleStem(nn.Module):
    def __init__(self, out_ch: int = 24):
        super().__init__()
        branch_ch = out_ch // 3
        rest = out_ch - branch_ch * 2
        self.b1 = ConvGNAct(1, branch_ch, kernel_size=7, stride=1)
        self.b2 = ConvGNAct(1, branch_ch, kernel_size=15, stride=1)
        self.b3 = ConvGNAct(1, rest, kernel_size=31, stride=1)
        self.mix = ConvGNAct(out_ch, out_ch, kernel_size=1, stride=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = torch.cat([self.b1(x), self.b2(x), self.b3(x)], dim=1)
        return self.mix(x)


class DispNetCNNV23(nn.Module):
    def __init__(
        self,
        input_length: int = 1536,
        base_channels: int = 24,
        output_dim: int = 49,
        dropout: float = 0.1,
        reference_disp: Optional[torch.Tensor] = None,
    ):
        super().__init__()
        self.input_length = input_length
        self.output_dim = output_dim
        if reference_disp is None:
            reference_disp = torch.zeros(output_dim, dtype=torch.float32)
        reference_disp = torch.as_tensor(reference_disp, dtype=torch.float32).reshape(output_dim)
        self.register_buffer("reference_disp", reference_disp)

        ch1 = base_channels
        ch2 = base_channels * 2
        ch3 = base_channels * 3
        ch4 = base_channels * 4
        ch5 = base_channels * 6

        self.stem = MultiScaleStem(ch1)
        self.backbone = nn.Sequential(
            ResidualBlock(ch1, ch1, stride=1, kernel_size=7),
            ResidualBlock(ch1, ch2, stride=2, kernel_size=9),
            ResidualBlock(ch2, ch2, stride=1, kernel_size=7, dilation=2),
            ResidualBlock(ch2, ch3, stride=2, kernel_size=9),
            ResidualBlock(ch3, ch3, stride=1, kernel_size=7, dilation=2),
            ResidualBlock(ch3, ch4, stride=2, kernel_size=9),
            ResidualBlock(ch4, ch4, stride=1, kernel_size=7, dilation=2),
            ResidualBlock(ch4, ch5, stride=2, kernel_size=9),
            ResidualBlock(ch5, ch5, stride=1, kernel_size=7, dilation=2),
            ResidualBlock(ch5, ch5, stride=2, kernel_size=7),
        )

        feat_dim = ch5 * 2
        hidden = max(base_channels * 8, 128)
        self.head = nn.Sequential(
            nn.Linear(feat_dim, hidden),
            nn.LayerNorm(hidden),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(hidden, hidden),
            nn.SiLU(inplace=True),
            nn.Dropout(dropout),
        )
        self.to_disp_delta = nn.Linear(hidden, output_dim)
        self.to_certainty_logits = nn.Linear(hidden, output_dim)
        nn.init.normal_(self.to_disp_delta.weight, mean=0.0, std=1e-3)
        nn.init.zeros_(self.to_disp_delta.bias)

    def set_reference_disp(self, reference_disp: torch.Tensor):
        reference_disp = torch.as_tensor(reference_disp, dtype=self.reference_disp.dtype, device=self.reference_disp.device)
        reference_disp = reference_disp.reshape(self.output_dim)
        self.reference_disp.copy_(reference_disp)

    @staticmethod
    def normalize_waveform(waveform: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
        mean = waveform.mean(dim=-1, keepdim=True)
        std = waveform.std(dim=-1, keepdim=True).clamp_min(eps)
        return (waveform - mean) / std

    def forward(self, waveform: torch.Tensor) -> Dict[str, torch.Tensor]:
        if waveform.ndim == 2:
            waveform = waveform.unsqueeze(1)
        elif waveform.ndim == 3 and waveform.size(1) != 1:
            raise ValueError(f"Expected waveform shape [B, T] or [B, 1, T], got {tuple(waveform.shape)}")

        waveform = self.normalize_waveform(waveform)
        h = self.backbone(self.stem(waveform))
        avg_pool = h.mean(dim=-1)
        max_pool = h.amax(dim=-1)
        feat = torch.cat([avg_pool, max_pool], dim=1)
        feat = self.head(feat)

        disp_delta = self.to_disp_delta(feat)
        disp_mu = self.reference_disp.view(1, -1) + disp_delta
        certainty_logits = self.to_certainty_logits(feat)
        return {
            "disp_mu": disp_mu,
            "disp_delta": disp_delta,
            "reference_disp": self.reference_disp.view(1, -1).expand_as(disp_mu),
            "certainty_logits": certainty_logits,
            "certainty": torch.sigmoid(certainty_logits),
        }


def masked_huber_loss(pred, target, mask, delta: float = 0.05, eps: float = 1e-8):
    abs_err = (pred - target).abs()
    huber = torch.where(abs_err <= delta, 0.5 * abs_err.pow(2) / delta, abs_err - 0.5 * delta)
    huber = huber * mask
    valid_count = mask.sum(dim=1)
    sample_valid = (valid_count > 0).float()
    per_sample = huber.sum(dim=1) / valid_count.clamp_min(eps)
    per_sample = per_sample * sample_valid
    return per_sample.sum() / sample_valid.sum().clamp_min(eps)


def masked_slope_loss(pred, target, mask, eps: float = 1e-8):
    pair_mask = mask[:, 1:] * mask[:, :-1]
    pred_slope = pred[:, 1:] - pred[:, :-1]
    target_slope = target[:, 1:] - target[:, :-1]
    per_elem = F.smooth_l1_loss(pred_slope, target_slope, reduction="none") * pair_mask
    valid_count = pair_mask.sum(dim=1)
    sample_valid = (valid_count > 0).float()
    per_sample = per_elem.sum(dim=1) / valid_count.clamp_min(eps)
    per_sample = per_sample * sample_valid
    return per_sample.sum() / sample_valid.sum().clamp_min(eps)


def masked_curvature_loss(pred, target, mask, eps: float = 1e-8):
    if pred.size(1) < 3:
        return pred.new_tensor(0.0)
    triplet_mask = mask[:, 2:] * mask[:, 1:-1] * mask[:, :-2]
    pred_d2 = pred[:, 2:] - 2.0 * pred[:, 1:-1] + pred[:, :-2]
    target_d2 = target[:, 2:] - 2.0 * target[:, 1:-1] + target[:, :-2]
    per_elem = F.smooth_l1_loss(pred_d2, target_d2, reduction="none") * triplet_mask
    valid_count = triplet_mask.sum(dim=1)
    sample_valid = (valid_count > 0).float()
    per_sample = per_elem.sum(dim=1) / valid_count.clamp_min(eps)
    per_sample = per_sample * sample_valid
    return per_sample.sum() / sample_valid.sum().clamp_min(eps)


def masked_pairwise_delta_loss(pred_delta, target_delta, mask, delta: float = 0.05, eps: float = 1e-8):
    if pred_delta.size(0) < 2:
        return pred_delta.new_tensor(0.0)
    pred_shift = torch.roll(pred_delta, shifts=1, dims=0)
    target_shift = torch.roll(target_delta, shifts=1, dims=0)
    mask_shift = torch.roll(mask, shifts=1, dims=0)
    pair_mask = mask * mask_shift
    if pair_mask.sum() <= 0:
        return pred_delta.new_tensor(0.0)

    pred_pair = pred_delta - pred_shift
    target_pair = target_delta - target_shift
    abs_err = (pred_pair - target_pair).abs()
    huber = torch.where(abs_err <= delta, 0.5 * abs_err.pow(2) / delta, abs_err - 0.5 * delta)
    return (huber * pair_mask).sum() / pair_mask.sum().clamp_min(eps)


def masked_period_std_loss(pred_delta, target_delta, mask, eps: float = 1e-8):
    valid_count = mask.sum(dim=0)
    usable = valid_count >= 2
    if not torch.any(usable):
        return pred_delta.new_tensor(0.0)

    count = valid_count.clamp_min(1.0)
    pred_mean = (pred_delta * mask).sum(dim=0) / count
    target_mean = (target_delta * mask).sum(dim=0) / count
    pred_var = (((pred_delta - pred_mean.view(1, -1)) * mask).pow(2).sum(dim=0) / count).clamp_min(0.0)
    target_var = (((target_delta - target_mean.view(1, -1)) * mask).pow(2).sum(dim=0) / count).clamp_min(0.0)
    pred_std = torch.sqrt(pred_var + eps)
    target_std = torch.sqrt(target_var + eps)
    return (pred_std[usable] - target_std[usable]).abs().mean()


def certainty_bce_loss(certainty_logits, target_mask, pos_weight=None):
    return F.binary_cross_entropy_with_logits(
        certainty_logits,
        target_mask,
        pos_weight=pos_weight,
        reduction="mean",
    )


def compute_total_loss(
    outputs: Dict[str, torch.Tensor],
    target_disp: torch.Tensor,
    target_mask: torch.Tensor,
    lambda_certainty: float = 0.2,
    lambda_slope: float = 0.25,
    lambda_curvature: float = 0.05,
    lambda_pairwise: float = 0.35,
    lambda_std: float = 0.75,
    huber_delta: float = 0.05,
    certainty_pos_weight: Optional[torch.Tensor] = None,
):
    pred = outputs["disp_mu"]
    pred_delta = outputs["disp_delta"]
    target_delta = target_disp - outputs["reference_disp"]
    rec_loss = masked_huber_loss(pred, target_disp, target_mask, delta=huber_delta)
    slope_loss = masked_slope_loss(pred_delta, target_delta, target_mask)
    curvature_loss = masked_curvature_loss(pred_delta, target_delta, target_mask)
    pairwise_loss = masked_pairwise_delta_loss(pred_delta, target_delta, target_mask, delta=huber_delta)
    std_loss = masked_period_std_loss(pred_delta, target_delta, target_mask)
    certainty_loss = certainty_bce_loss(outputs["certainty_logits"], target_mask, certainty_pos_weight)

    total = (
        rec_loss
        + lambda_certainty * certainty_loss
        + lambda_slope * slope_loss
        + lambda_curvature * curvature_loss
        + lambda_pairwise * pairwise_loss
        + lambda_std * std_loss
    )
    return {
        "loss": total,
        "rec_loss": rec_loss,
        "certainty_loss": certainty_loss,
        "slope_loss": slope_loss,
        "curvature_loss": curvature_loss,
        "pairwise_loss": pairwise_loss,
        "std_loss": std_loss,
    }


@torch.no_grad()
def masked_mae(pred, target, mask, eps: float = 1e-8):
    err = (pred - target).abs() * mask
    return err.sum() / mask.sum().clamp_min(eps)


@torch.no_grad()
def masked_rmse(pred, target, mask, eps: float = 1e-8):
    err2 = (pred - target).pow(2) * mask
    return torch.sqrt(err2.sum() / mask.sum().clamp_min(eps))


@torch.no_grad()
def certainty_f1_from_logits(certainty_logits, target_mask, threshold: float = 0.5, eps: float = 1e-8):
    pred = (torch.sigmoid(certainty_logits) > threshold).float()
    tp = (pred * target_mask).sum()
    fp = (pred * (1.0 - target_mask)).sum()
    fn = ((1.0 - pred) * target_mask).sum()
    precision = tp / (tp + fp + eps)
    recall = tp / (tp + fn + eps)
    return 2.0 * precision * recall / (precision + recall + eps)


def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device):
    return {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()}


def unwrap_waveform(waveform: torch.Tensor) -> torch.Tensor:
    if waveform.ndim == 3 and waveform.size(1) == 1:
        return waveform.squeeze(1)
    return waveform


def run_one_epoch(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    optimizer: Optional[torch.optim.Optimizer] = None,
    scaler: Optional[torch.cuda.amp.GradScaler] = None,
    use_amp: bool = True,
    grad_clip: Optional[float] = 1.0,
    lambda_certainty: float = 0.2,
    lambda_slope: float = 0.25,
    lambda_curvature: float = 0.05,
    lambda_pairwise: float = 0.35,
    lambda_std: float = 0.75,
    huber_delta: float = 0.05,
    certainty_pos_weight: Optional[torch.Tensor] = None,
):
    is_train = optimizer is not None
    model.train(is_train)

    totals = {
        "loss": 0.0,
        "rec_loss": 0.0,
        "certainty_loss": 0.0,
        "slope_loss": 0.0,
        "curvature_loss": 0.0,
        "pairwise_loss": 0.0,
        "std_loss": 0.0,
        "mae": 0.0,
        "rmse": 0.0,
        "certainty_f1": 0.0,
    }
    n_batches = 0

    for batch in loader:
        batch = move_batch_to_device(batch, device)
        waveform = unwrap_waveform(batch["waveform"].float())
        disp = batch["disp"].float()
        mask = batch["mask"].float()

        if is_train:
            optimizer.zero_grad(set_to_none=True)

        with torch.set_grad_enabled(is_train):
            with torch.cuda.amp.autocast(enabled=(use_amp and device.type == "cuda")):
                outputs = model(waveform)
                loss_dict = compute_total_loss(
                    outputs=outputs,
                    target_disp=disp,
                    target_mask=mask,
                    lambda_certainty=lambda_certainty,
                    lambda_slope=lambda_slope,
                    lambda_curvature=lambda_curvature,
                    lambda_pairwise=lambda_pairwise,
                    lambda_std=lambda_std,
                    huber_delta=huber_delta,
                    certainty_pos_weight=certainty_pos_weight,
                )
                loss = loss_dict["loss"]

            if is_train:
                if use_amp and device.type == "cuda":
                    scaler.scale(loss).backward()
                    if exists(grad_clip):
                        scaler.unscale_(optimizer)
                        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    loss.backward()
                    if exists(grad_clip):
                        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                    optimizer.step()

        with torch.no_grad():
            totals["mae"] += masked_mae(outputs["disp_mu"], disp, mask).item()
            totals["rmse"] += masked_rmse(outputs["disp_mu"], disp, mask).item()
            totals["certainty_f1"] += certainty_f1_from_logits(outputs["certainty_logits"], mask).item()

        for key in ["loss", "rec_loss", "certainty_loss", "slope_loss", "curvature_loss", "pairwise_loss", "std_loss"]:
            totals[key] += loss_dict[key].item()
        n_batches += 1

    denom = max(n_batches, 1)
    return {k: v / denom for k, v in totals.items()}


@torch.no_grad()
def estimate_reference_disp(loader: DataLoader, output_dim: int, device: torch.device):
    sums = torch.zeros(output_dim, device=device)
    counts = torch.zeros(output_dim, device=device)
    period_sum = torch.zeros(output_dim, device=device)

    for batch in loader:
        disp = batch["disp"].to(device).float()
        mask = batch["mask"].to(device).float()
        periods = batch["periods"].to(device).float()
        sums += (disp * mask).sum(dim=0)
        period_sum += (periods * mask).sum(dim=0)
        counts += mask.sum(dim=0)

    reference_disp = sums / counts.clamp_min(1.0)
    reference_disp = torch.where(counts > 0, reference_disp, torch.zeros_like(reference_disp))
    reference_periods = period_sum / counts.clamp_min(1.0)
    return reference_disp.detach().cpu(), reference_periods.detach().cpu(), counts.detach().cpu()


@dataclass
class TrainConfig:
    h5_path: str = "data/ncf_disp_dataset_with_disp_image.h5"
    save_dir: str = "ckpt_large/checkpoints_dispnet_v2.3_residual_cnn"

    waveform_length: int = 1536
    batch_size: int = 32
    num_workers: int = 4

    base_channels: int = 24
    output_dim: int = 49
    dropout: float = 0.1

    epochs: int = 80
    lr: float = 2e-4
    weight_decay: float = 1e-4
    lambda_certainty: float = 0.2
    lambda_slope: float = 0.25
    lambda_curvature: float = 0.05
    lambda_pairwise: float = 0.35
    lambda_std: float = 0.75
    huber_delta: float = 0.05
    grad_clip: float = 1.0
    use_amp: bool = True
    seed: int = 42
    min_lr: float = 1e-6
    use_certainty_pos_weight: bool = False
    device: str = default_device()


def save_checkpoint(path, model, optimizer, scheduler, epoch, best_val_loss, cfg: TrainConfig):
    torch.save(
        {
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict() if scheduler is not None else None,
            "epoch": epoch,
            "best_val_loss": best_val_loss,
            "config": cfg.__dict__,
            "model_name": "DispNetCNNV23",
            "reference_disp": model.reference_disp.detach().cpu(),
        },
        path,
    )


def main():
    from utils.dispdataset1d import build_dataloader

    cfg = TrainConfig()
    os.makedirs(cfg.save_dir, exist_ok=True)
    set_seed(cfg.seed)
    device = torch.device(cfg.device)

    train_loader = build_dataloader(
        h5_path=cfg.h5_path,
        split="train",
        batch_size=cfg.batch_size,
        num_workers=cfg.num_workers,
        waveform_length=cfg.waveform_length,
        random_ncf=True,
        pin_memory=(device.type == "cuda"),
        drop_last=False,
        seed=cfg.seed,
    )
    val_loader = build_dataloader(
        h5_path=cfg.h5_path,
        split="test",
        batch_size=cfg.batch_size,
        num_workers=cfg.num_workers,
        waveform_length=cfg.waveform_length,
        random_ncf=False,
        pin_memory=(device.type == "cuda"),
        drop_last=False,
        seed=cfg.seed,
    )

    print("[Info] estimating train-set reference dispersion...")
    reference_disp, reference_periods, reference_counts = estimate_reference_disp(
        train_loader,
        output_dim=cfg.output_dim,
        device=device,
    )
    print(
        f"[Info] reference dispersion: "
        f"valid_periods={(reference_counts > 0).sum().item()}/{cfg.output_dim}, "
        f"v_min={reference_disp[reference_counts > 0].min().item():.4f}, "
        f"v_max={reference_disp[reference_counts > 0].max().item():.4f}"
    )

    model = DispNetCNNV23(
        input_length=cfg.waveform_length,
        base_channels=cfg.base_channels,
        output_dim=cfg.output_dim,
        dropout=cfg.dropout,
        reference_disp=reference_disp,
    ).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.epochs, eta_min=cfg.min_lr)
    scaler = torch.cuda.amp.GradScaler(enabled=(cfg.use_amp and device.type == "cuda"))

    certainty_pos_weight = None
    if cfg.use_certainty_pos_weight:
        certainty_pos_weight = torch.tensor([2.0], device=device)

    best_val_loss = float("inf")
    print("========== DispNet v2.3 Residual CNN Training Start ==========")
    print(f"device        : {device}")
    print(f"params        : {sum(p.numel() for p in model.parameters())}")
    print(f"train batches : {len(train_loader)}")
    print(f"val batches   : {len(val_loader)}")
    print(f"save_dir      : {cfg.save_dir}")

    for epoch in range(1, cfg.epochs + 1):
        t0 = time.time()
        train_stats = run_one_epoch(
            model=model,
            loader=train_loader,
            device=device,
            optimizer=optimizer,
            scaler=scaler,
            use_amp=cfg.use_amp,
            grad_clip=cfg.grad_clip,
            lambda_certainty=cfg.lambda_certainty,
            lambda_slope=cfg.lambda_slope,
            lambda_curvature=cfg.lambda_curvature,
            lambda_pairwise=cfg.lambda_pairwise,
            lambda_std=cfg.lambda_std,
            huber_delta=cfg.huber_delta,
            certainty_pos_weight=certainty_pos_weight,
        )
        val_stats = run_one_epoch(
            model=model,
            loader=val_loader,
            device=device,
            optimizer=None,
            scaler=None,
            use_amp=cfg.use_amp,
            grad_clip=None,
            lambda_certainty=cfg.lambda_certainty,
            lambda_slope=cfg.lambda_slope,
            lambda_curvature=cfg.lambda_curvature,
            lambda_pairwise=cfg.lambda_pairwise,
            lambda_std=cfg.lambda_std,
            huber_delta=cfg.huber_delta,
            certainty_pos_weight=certainty_pos_weight,
        )
        scheduler.step()
        dt = time.time() - t0
        lr_now = optimizer.param_groups[0]["lr"]

        print(
            f"[Epoch {epoch:03d}/{cfg.epochs:03d}] "
            f"time={dt:.1f}s lr={lr_now:.2e} | "
            f"train: loss={train_stats['loss']:.5f}, rec={train_stats['rec_loss']:.5f}, "
            f"cert={train_stats['certainty_loss']:.5f}, slope={train_stats['slope_loss']:.5f}, "
            f"curv={train_stats['curvature_loss']:.5f}, pair={train_stats['pairwise_loss']:.5f}, "
            f"std={train_stats['std_loss']:.5f}, mae={train_stats['mae']:.5f}, "
            f"rmse={train_stats['rmse']:.5f}, cert_f1={train_stats['certainty_f1']:.5f} | "
            f"val: loss={val_stats['loss']:.5f}, rec={val_stats['rec_loss']:.5f}, "
            f"cert={val_stats['certainty_loss']:.5f}, slope={val_stats['slope_loss']:.5f}, "
            f"curv={val_stats['curvature_loss']:.5f}, pair={val_stats['pairwise_loss']:.5f}, "
            f"std={val_stats['std_loss']:.5f}, mae={val_stats['mae']:.5f}, "
            f"rmse={val_stats['rmse']:.5f}, cert_f1={val_stats['certainty_f1']:.5f}"
        )

        latest_path = os.path.join(cfg.save_dir, "latest.pt")
        save_checkpoint(latest_path, model, optimizer, scheduler, epoch, best_val_loss, cfg)

        if val_stats["loss"] < best_val_loss:
            best_val_loss = val_stats["loss"]
            best_path = os.path.join(cfg.save_dir, "best.pt")
            save_checkpoint(best_path, model, optimizer, scheduler, epoch, best_val_loss, cfg)
            print(f"  -> Best model saved to {best_path}")

    print("========== DispNet v2.3 Residual CNN Training Done ==========")
    print(f"Best val loss: {best_val_loss:.6f}")


if __name__ == "__main__":
    main()