Building/src_new/train.py

from __future__ import annotations

import argparse
import random
from dataclasses import asdict
from pathlib import Path
from typing import Any

import pandas as pd
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import ReduceLROnPlateau

try:
    from .config import ExperimentConfig, make_rms_forward_config
    from .dataset import build_dataloaders, report_to_text
    from .model import build_model, compute_receptive_field
except ImportError:
    from config import ExperimentConfig, make_rms_forward_config
    from dataset import build_dataloaders, report_to_text
    from model import build_model, compute_receptive_field


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train TCN for direct RMS regression.")
    parser.add_argument("--epochs", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--device", type=str, default=None)
    return parser.parse_args()


def set_seed(seed: int) -> None:
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def resolve_device(device_name: str) -> torch.device:
    if device_name.startswith("cuda") and not torch.cuda.is_available():
        return torch.device("cpu")
    return torch.device(device_name)


def serialize_for_checkpoint(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {key: serialize_for_checkpoint(sub_value) for key, sub_value in value.items()}
    if isinstance(value, tuple):
        return [serialize_for_checkpoint(item) for item in value]
    if isinstance(value, list):
        return [serialize_for_checkpoint(item) for item in value]
    return value


def denormalize_target(
    pred_norm: torch.Tensor,
    target_norm: torch.Tensor,
    target_mean: torch.Tensor,
    target_std: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    pred = pred_norm * target_std + target_mean
    target = target_norm * target_std + target_mean
    return pred, target


def relative_rms_error(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return torch.abs(pred - target) / torch.clamp(target.abs(), min=1e-6)


class RMSRegressionLoss(nn.Module):
    def __init__(
        self,
        relative_rms_weight: float,
        log_rms_weight: float,
        mae_weight: float,
        waveform_l1_weight: float,
        waveform_huber_weight: float,
    ) -> None:
        super().__init__()
        self.relative_rms_weight = relative_rms_weight
        self.log_rms_weight = log_rms_weight
        self.mae_weight = mae_weight
        self.waveform_l1_weight = waveform_l1_weight
        self.waveform_huber_weight = waveform_huber_weight

    def forward(
        self,
        pred: torch.Tensor,
        target: torch.Tensor,
        pred_wave: torch.Tensor,
        target_wave: torch.Tensor,
        mask: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        rel_loss = relative_rms_error(pred, target).mean()
        log_loss = F.huber_loss(torch.log(torch.clamp(pred, min=1e-6)), torch.log(torch.clamp(target, min=1e-6)))
        mae_loss = F.l1_loss(pred, target)
        mask_expanded = mask.unsqueeze(-1)
        valid_count = mask_expanded.sum().clamp_min(1.0)
        wave_residual = (pred_wave - target_wave) * mask_expanded
        waveform_l1 = torch.abs(wave_residual).sum() / valid_count
        waveform_huber = F.huber_loss(pred_wave * mask_expanded, target_wave * mask_expanded, reduction="sum") / valid_count
        total = (
            self.relative_rms_weight * rel_loss
            + self.log_rms_weight * log_loss
            + self.mae_weight * mae_loss
            + self.waveform_l1_weight * waveform_l1
            + self.waveform_huber_weight * waveform_huber
        )
        return {
            "total": total,
            "relative": rel_loss,
            "log": log_loss,
            "mae": mae_loss,
            "waveform_l1": waveform_l1,
            "waveform_huber": waveform_huber,
        }


def run_epoch(
    model: nn.Module,
    dataloader: torch.utils.data.DataLoader,
    optimizer: AdamW | None,
    criterion: RMSRegressionLoss,
    target_mean: torch.Tensor,
    target_std: torch.Tensor,
    device: torch.device,
    grad_clip_norm: float,
    scaler: torch.cuda.amp.GradScaler,
    amp_enabled: bool,
) -> dict[str, float]:
    is_train = optimizer is not None
    model.train(is_train)

    total_loss_sum = 0.0
    relative_loss_sum = 0.0
    log_loss_sum = 0.0
    mae_loss_sum = 0.0
    waveform_l1_sum = 0.0
    waveform_huber_sum = 0.0
    rms_error_sum = 0.0
    sample_count = 0

    for batch in dataloader:
        x = batch["x"].to(device)
        aux = batch["aux"].to(device)
        mask = batch["mask"].to(device)
        y_norm = batch["y"].to(device)
        y_wave = batch["y_wave"].to(device)
        x_rms_raw = batch["x_rms_raw"].to(device)
        y_rms_raw = batch["y_raw"].to(device)

        if is_train:
            optimizer.zero_grad(set_to_none=True)

        with torch.amp.autocast(device_type=device.type, enabled=amp_enabled):
            outputs = model(x, mask, aux)
            pred_target, _ = denormalize_target(outputs["rms"], y_norm, target_mean, target_std)
            pred = torch.exp(pred_target) * x_rms_raw
            target = y_rms_raw
            losses = criterion(pred, target, outputs["waveform"], y_wave, mask)

        if is_train:
            scaler.scale(losses["total"]).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip_norm)
            scaler.step(optimizer)
            scaler.update()

        batch_size = x.shape[0]
        total_loss_sum += losses["total"].detach().item() * batch_size
        relative_loss_sum += losses["relative"].detach().item() * batch_size
        log_loss_sum += losses["log"].detach().item() * batch_size
        mae_loss_sum += losses["mae"].detach().item() * batch_size
        waveform_l1_sum += losses["waveform_l1"].detach().item() * batch_size
        waveform_huber_sum += losses["waveform_huber"].detach().item() * batch_size
        rms_error_sum += relative_rms_error(pred.detach(), target.detach()).mean().item() * batch_size
        sample_count += batch_size

    return {
        "loss": total_loss_sum / sample_count,
        "relative_loss": relative_loss_sum / sample_count,
        "log_loss": log_loss_sum / sample_count,
        "mae_loss": mae_loss_sum / sample_count,
        "waveform_l1": waveform_l1_sum / sample_count,
        "waveform_huber": waveform_huber_sum / sample_count,
        "rms_error": rms_error_sum / sample_count,
    }


def checkpoint_paths(config: ExperimentConfig) -> tuple[Path, Path]:
    root = config.data.project_root / config.train.checkpoint_dir / "forward_rms"
    root.mkdir(parents=True, exist_ok=True)
    return root / config.train.best_model_name, root / config.train.history_name

def train_model(config: ExperimentConfig) -> None:
    set_seed(config.train.seed)
    device = resolve_device(config.train.device)
    loaders, datasets, reports = build_dataloaders(config)
    train_dataset = datasets["train"]
    normalization = train_dataset.normalization
    target_mean = normalization.target_mean.to(device).view(1, 1)
    target_std = normalization.target_std.to(device).view(1, 1)

    model = build_model(config).to(device)
    receptive_field = compute_receptive_field(config.model)
    optimizer = AdamW(model.parameters(), lr=config.train.learning_rate, weight_decay=config.train.weight_decay)
    scheduler = ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=config.train.lr_scheduler_factor,
        patience=config.train.lr_scheduler_patience,
        min_lr=config.train.min_learning_rate,
    )
    criterion = RMSRegressionLoss(
        relative_rms_weight=config.loss.relative_rms_weight,
        log_rms_weight=config.loss.log_rms_weight,
        mae_weight=config.loss.mae_weight,
        waveform_l1_weight=config.loss.waveform_l1_weight,
        waveform_huber_weight=config.loss.waveform_huber_weight,
    )
    amp_enabled = config.train.use_amp and device.type == "cuda"
    scaler = torch.cuda.amp.GradScaler(enabled=amp_enabled)

    best_val_error = float("inf")
    epochs_without_improvement = 0
    history: list[dict[str, float]] = []
    best_model_path, history_path = checkpoint_paths(config)

    print(f"Device: {device}")
    print(f"TCN receptive field: {receptive_field.receptive_field} samples, dilations={receptive_field.dilations}")
    print(report_to_text(reports))

    for epoch in range(1, config.train.epochs + 1):
        train_metrics = run_epoch(
            model=model,
            dataloader=loaders["train"],
            optimizer=optimizer,
            criterion=criterion,
            target_mean=target_mean,
            target_std=target_std,
            device=device,
            grad_clip_norm=config.train.grad_clip_norm,
            scaler=scaler,
            amp_enabled=amp_enabled,
        )
        val_metrics = run_epoch(
            model=model,
            dataloader=loaders["val"],
            optimizer=None,
            criterion=criterion,
            target_mean=target_mean,
            target_std=target_std,
            device=device,
            grad_clip_norm=config.train.grad_clip_norm,
            scaler=scaler,
            amp_enabled=amp_enabled,
        )
        scheduler.step(val_metrics["rms_error"])
        current_lr = optimizer.param_groups[0]["lr"]

        history_row = {
            "epoch": epoch,
            "lr": current_lr,
            "train_loss": train_metrics["loss"],
            "train_relative_loss": train_metrics["relative_loss"],
            "train_log_loss": train_metrics["log_loss"],
            "train_mae_loss": train_metrics["mae_loss"],
            "train_waveform_l1": train_metrics["waveform_l1"],
            "train_waveform_huber": train_metrics["waveform_huber"],
            "train_rms_error": train_metrics["rms_error"],
            "val_loss": val_metrics["loss"],
            "val_relative_loss": val_metrics["relative_loss"],
            "val_log_loss": val_metrics["log_loss"],
            "val_mae_loss": val_metrics["mae_loss"],
            "val_waveform_l1": val_metrics["waveform_l1"],
            "val_waveform_huber": val_metrics["waveform_huber"],
            "val_rms_error": val_metrics["rms_error"],
        }
        history.append(history_row)

        print(
            f"Epoch {epoch:03d} | train_loss={train_metrics['loss']:.6f} | "
            f"val_loss={val_metrics['loss']:.6f} | val_rms_error={val_metrics['rms_error']:.6f} | lr={current_lr:.2e}"
        )

        if val_metrics["rms_error"] < best_val_error:
            best_val_error = val_metrics["rms_error"]
            epochs_without_improvement = 0
            torch.save(
                {
                    "epoch": epoch,
                    "model_state_dict": model.state_dict(),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "best_val_rms_error": best_val_error,
                    "config": serialize_for_checkpoint(asdict(config)),
                },
                best_model_path,
            )
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= config.train.early_stop_patience:
                print(f"Early stopping triggered after {epoch} epochs.")
                break

    history_df = pd.DataFrame(history)
    history_df.to_csv(history_path, index=False)
    print(f"Best model saved to: {best_model_path}")
    print(f"Training history saved to: {history_path}")


def main() -> None:
    args = parse_args()
    config = make_rms_forward_config()
    if args.epochs is not None:
        config.train.epochs = args.epochs
    if args.batch_size is not None:
        config.data.batch_size = args.batch_size
    if args.device is not None:
        config.train.device = args.device
    train_model(config)


if __name__ == "__main__":
    main()