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()