#!/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()