from __future__ import annotations import json import random from dataclasses import dataclass from pathlib import Path import joblib import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset from src.models.forward_surrogate import ForwardSurrogate """正演代理模型训练循环。 训练目标是拼接后的曲线向量：log_pressure、log_derivative 和辅助 slope。 默认损失主要约束压力与导数曲线，slope 通常只记录监控，权重为 0。 """ class ForwardDataset(Dataset): """将预处理后的 numpy 数组封装为 PyTorch Dataset。""" def __init__(self, params_x: np.ndarray, schedule_x: np.ndarray, curve_y: np.ndarray): self.params_x = torch.tensor(params_x, dtype=torch.float32) self.schedule_x = torch.tensor(schedule_x, dtype=torch.float32) self.curve_y = torch.tensor(curve_y, dtype=torch.float32) def __len__(self) -> int: return len(self.params_x) def __getitem__(self, idx: int): return self.params_x[idx], self.schedule_x[idx], self.curve_y[idx] @dataclass class TrainConfig: """scripts/train_forward.py 使用的训练运行参数。""" processed_path: Path output_dir: Path seed: int = 42 batch_size: int = 128 epochs: int = 100 lr: float = 1e-3 weight_decay: float = 1e-5 hidden_dim: int = 128 dropout: float = 0.0 w_pressure: float = 1.0 w_derivative: float = 2.0 w_slope: float = 0.0 w_bias_pressure: float = 0.15 w_bias_derivative: float = 0.05 w_derivative_shape: float = 0.10 w_autofit_pressure: float = 0.0 w_autofit_derivative: float = 0.0 use_huber: bool = True huber_beta: float = 0.05 use_sample_reweight: bool = True sample_reweight_alpha: float = 0.4 sample_weight_min: float = 1.0 sample_weight_max: float = 2.5 use_schedule: bool = True device: str = "cuda" if torch.cuda.is_available() else "cpu" def set_global_seed(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def load_processed_dataset(path: Path) -> dict: """加载预处理后的 joblib 数据，并检查必需的数据划分字段。""" data = joblib.load(path) required_keys = [ "X_params_train", "X_schedule_train", "Y_curve_train", "X_params_val", "X_schedule_val", "Y_curve_val", "X_params_test", "X_schedule_test", "Y_curve_test", ] for k in required_keys: if k not in data: raise KeyError(f"processed dataset 缺少字段: {k}") return data def infer_curve_layout(data: dict) -> dict: """从元数据读取曲线布局；没有元数据时回退为旧的三段布局。""" meta = data.get("meta", {}) curve_dim = int(meta.get("curve_dim", data["Y_curve_train"].shape[1])) curve_layout = meta.get("curve_layout") if curve_layout is not None: return curve_layout if curve_dim % 3 != 0: raise ValueError(f"curve_dim={curve_dim} 不能按 3 段均分，且 meta 中没有 curve_layout") n_time_points = curve_dim // 3 return { "n_time_points": int(n_time_points), "parts": [ {"name": "log_pressure", "start": 0, "end": n_time_points}, {"name": "log_derivative", "start": n_time_points, "end": 2 * n_time_points}, {"name": "slope", "start": 2 * n_time_points, "end": 3 * n_time_points}, ], } def get_part_slices(curve_layout: dict) -> dict[str, slice]: """将 curve_layout 元数据转换为 pressure/derivative/slope 的切片。""" out: dict[str, slice] = {} for part in curve_layout["parts"]: name = str(part["name"]) start = int(part["start"]) end = int(part["end"]) out[name] = slice(start, end) return out def smooth_l1_per_sample(pred: torch.Tensor, target: torch.Tensor, beta: float) -> torch.Tensor: diff = torch.abs(pred - target) loss = torch.where( diff < beta, 0.5 * diff * diff / beta, diff - 0.5 * beta, ) return loss.mean(dim=1) def mse_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: return ((pred - target) ** 2).mean(dim=1) def l1_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: return torch.abs(pred - target).mean(dim=1) def first_diff(x: torch.Tensor) -> torch.Tensor: return x[:, 1:] - x[:, :-1] def affine_restore(x_scaled: torch.Tensor, mean: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: """在 torch 中还原 StandardScaler 变换，用于计算自动拟合风格损失。""" return x_scaled * scale + mean def autofit_curve_objective_per_sample(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """混合相对误差/绝对误差曲线目标函数的 torch 版本。""" weight_factor = torch.clamp(torch.abs(target) * 0.01, max=100.0) weight = 1.0 / (1.0 + weight_factor) scale = torch.maximum(torch.maximum(torch.abs(target), torch.abs(pred)), torch.full_like(target, 1e-12)) relative_error = torch.abs(target - pred) / scale absolute_error = torch.abs(target - pred) point_error = 0.7 * relative_error + 0.3 * absolute_error weighted_mse = (weight * (point_error**2)).sum(dim=1) / torch.clamp(weight.sum(dim=1), min=1e-12) return torch.sqrt(weighted_mse) def build_sample_weight( true_p: torch.Tensor, true_d: torch.Tensor, alpha: float, w_min: float, w_max: float, ) -> torch.Tensor: """训练时给高幅值曲线样本适度更高的权重。""" p_level = true_p.abs().mean(dim=1) d_level = true_d.abs().mean(dim=1) p_norm = p_level / (p_level.mean().detach() + 1e-6) d_norm = d_level / (d_level.mean().detach() + 1e-6) raw = 0.5 * p_norm + 0.5 * d_norm weight = 1.0 + alpha * (raw - 1.0) weight = torch.clamp(weight, min=w_min, max=w_max) return weight def compute_weighted_loss( pred: torch.Tensor, target: torch.Tensor, slices: dict[str, slice], curve_mean_raw: torch.Tensor, curve_scale_raw: torch.Tensor, huber_beta: float, w_pressure: float, w_derivative: float, w_slope: float, w_bias_pressure: float, w_bias_derivative: float, w_derivative_shape: float, w_autofit_pressure: float, w_autofit_derivative: float, use_sample_reweight: bool, sample_reweight_alpha: float, sample_weight_min: float, sample_weight_max: float, ): """计算正演代理模型训练和监控使用的所有损失分量。""" pred_p = pred[:, slices["log_pressure"]] pred_d = pred[:, slices["log_derivative"]] pred_s = pred[:, slices["slope"]] true_p = target[:, slices["log_pressure"]] true_d = target[:, slices["log_derivative"]] true_s = target[:, slices["slope"]] mean_p = curve_mean_raw[slices["log_pressure"]].unsqueeze(0) scale_p = curve_scale_raw[slices["log_pressure"]].unsqueeze(0) mean_d = curve_mean_raw[slices["log_derivative"]].unsqueeze(0) scale_d = curve_scale_raw[slices["log_derivative"]].unsqueeze(0) loss_p_vec = smooth_l1_per_sample(pred_p, true_p, beta=huber_beta) loss_d_vec = smooth_l1_per_sample(pred_d, true_d, beta=huber_beta) loss_s_vec = mse_per_sample(pred_s, true_s) # 偏置损失用于抑制整条预测曲线的纵向漂移； # 点误差负责学习更细的局部形态。 pred_p_mean = pred_p.mean(dim=1, keepdim=True) true_p_mean = true_p.mean(dim=1, keepdim=True) loss_bias_p_vec = l1_per_sample(pred_p_mean, true_p_mean) pred_d_mean = pred_d.mean(dim=1, keepdim=True) true_d_mean = true_d.mean(dim=1, keepdim=True) loss_bias_d_vec = l1_per_sample(pred_d_mean, true_d_mean) pred_d_diff = first_diff(pred_d) true_d_diff = first_diff(true_d) loss_d_shape_vec = smooth_l1_per_sample(pred_d_diff, true_d_diff, beta=huber_beta) # 自动拟合风格损失在原始曲线尺度上计算。默认权重为 0， # 但保留这个接口，方便后续实验启用。 pred_p_raw = affine_restore(pred_p, mean_p, scale_p) true_p_raw = affine_restore(true_p, mean_p, scale_p) pred_d_raw = affine_restore(pred_d, mean_d, scale_d) true_d_raw = affine_restore(true_d, mean_d, scale_d) loss_autofit_p_vec = autofit_curve_objective_per_sample(pred_p_raw, true_p_raw) loss_autofit_d_vec = autofit_curve_objective_per_sample(pred_d_raw, true_d_raw) total_vec = ( w_pressure * loss_p_vec + w_derivative * loss_d_vec + w_slope * loss_s_vec + w_bias_pressure * loss_bias_p_vec + w_bias_derivative * loss_bias_d_vec + w_derivative_shape * loss_d_shape_vec + w_autofit_pressure * loss_autofit_p_vec + w_autofit_derivative * loss_autofit_d_vec ) if use_sample_reweight: sample_weight = build_sample_weight( true_p=true_p, true_d=true_d, alpha=sample_reweight_alpha, w_min=sample_weight_min, w_max=sample_weight_max, ) else: sample_weight = torch.ones_like(total_vec) total = (total_vec * sample_weight).mean() return { "loss": total, "loss_pressure": loss_p_vec.mean(), "loss_derivative": loss_d_vec.mean(), "loss_slope": loss_s_vec.mean(), "loss_bias_pressure": loss_bias_p_vec.mean(), "loss_bias_derivative": loss_bias_d_vec.mean(), "loss_derivative_shape": loss_d_shape_vec.mean(), "loss_autofit_pressure": loss_autofit_p_vec.mean(), "loss_autofit_derivative": loss_autofit_d_vec.mean(), "sample_weight_mean": sample_weight.mean(), "sample_weight_max": sample_weight.max(), } def model_forward(model: nn.Module, params_x: torch.Tensor, schedule_x: torch.Tensor, use_schedule: bool) -> torch.Tensor: if use_schedule: return model(params_x, schedule_x) return model(params_x, None) def evaluate( model: nn.Module, loader: DataLoader, device: str, slices: dict[str, slice], cfg: TrainConfig, ) -> dict: model.eval() total = { "loss": 0.0, "loss_pressure": 0.0, "loss_derivative": 0.0, "loss_slope": 0.0, "loss_bias_pressure": 0.0, "loss_bias_derivative": 0.0, "loss_derivative_shape": 0.0, "loss_autofit_pressure": 0.0, "loss_autofit_derivative": 0.0, "sample_weight_mean": 0.0, "sample_weight_max": 0.0, } total_n = 0 with torch.no_grad(): for params_x, schedule_x, curve_y in loader: params_x = params_x.to(device) schedule_x = schedule_x.to(device) curve_y = curve_y.to(device) pred = model_forward(model, params_x, schedule_x, cfg.use_schedule) losses = compute_weighted_loss( pred=pred, target=curve_y, slices=slices, curve_mean_raw=cfg.curve_mean_raw, curve_scale_raw=cfg.curve_scale_raw, huber_beta=cfg.huber_beta, w_pressure=cfg.w_pressure, w_derivative=cfg.w_derivative, w_slope=cfg.w_slope, w_bias_pressure=cfg.w_bias_pressure, w_bias_derivative=cfg.w_bias_derivative, w_derivative_shape=cfg.w_derivative_shape, w_autofit_pressure=cfg.w_autofit_pressure, w_autofit_derivative=cfg.w_autofit_derivative, use_sample_reweight=cfg.use_sample_reweight, sample_reweight_alpha=cfg.sample_reweight_alpha, sample_weight_min=cfg.sample_weight_min, sample_weight_max=cfg.sample_weight_max, ) bs = params_x.size(0) for k in total: total[k] += losses[k].item() * bs total_n += bs denom = max(total_n, 1) return {k: v / denom for k, v in total.items()} def train_forward(cfg: TrainConfig) -> None: """训练正演代理模型，保存最优 checkpoint，并在测试集上评估。""" cfg.output_dir.mkdir(parents=True, exist_ok=True) set_global_seed(int(cfg.seed)) data = load_processed_dataset(cfg.processed_path) curve_layout = infer_curve_layout(data) part_slices = get_part_slices(curve_layout) scaler_curve = data["scaler_curve"] curve_mean_raw = np.asarray(scaler_curve.mean_, dtype=np.float32).reshape(-1) curve_scale_raw = np.asarray(scaler_curve.scale_, dtype=np.float32).reshape(-1) # 将曲线 scaler 张量挂到 cfg 上，方便训练和评估共用同一套损失代码。 cfg.curve_mean_raw = torch.tensor(curve_mean_raw, dtype=torch.float32, device=cfg.device) cfg.curve_scale_raw = torch.tensor(curve_scale_raw, dtype=torch.float32, device=cfg.device) train_ds = ForwardDataset( data["X_params_train"], data["X_schedule_train"], data["Y_curve_train"] ) val_ds = ForwardDataset( data["X_params_val"], data["X_schedule_val"], data["Y_curve_val"] ) test_ds = ForwardDataset( data["X_params_test"], data["X_schedule_test"], data["Y_curve_test"] ) loader_generator = torch.Generator() loader_generator.manual_seed(int(cfg.seed)) train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, generator=loader_generator) val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False) test_loader = DataLoader(test_ds, batch_size=cfg.batch_size, shuffle=False) param_dim = data["X_params_train"].shape[1] schedule_dim = data["X_schedule_train"].shape[1] curve_dim = data["Y_curve_train"].shape[1] # 模型维度从预处理数据集中自动推断，便于不同数据集版本共用训练入口。 model = ForwardSurrogate( param_dim=param_dim, schedule_dim=schedule_dim, curve_dim=curve_dim, hidden_dim=cfg.hidden_dim, dropout=cfg.dropout, use_schedule=cfg.use_schedule, ).to(cfg.device) optimizer = torch.optim.Adam( model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay, ) best_val = float("inf") best_path = cfg.output_dir / "forward_surrogate_best.pt" history: list[dict] = [] print("训练配置:") print(f" device={cfg.device}") print(f" seed={cfg.seed}") print(f" batch_size={cfg.batch_size}, epochs={cfg.epochs}, lr={cfg.lr}, weight_decay={cfg.weight_decay}") print(f" hidden_dim={cfg.hidden_dim}, dropout={cfg.dropout}") print(f" use_schedule={cfg.use_schedule}") print( f" weights: pressure={cfg.w_pressure}, derivative={cfg.w_derivative}, " f"slope={cfg.w_slope}, bias_p={cfg.w_bias_pressure}, " f"bias_d={cfg.w_bias_derivative}, d_shape={cfg.w_derivative_shape}, " f"autofit_p={cfg.w_autofit_pressure}, autofit_d={cfg.w_autofit_derivative}" ) print( f" sample_reweight={cfg.use_sample_reweight}, alpha={cfg.sample_reweight_alpha}, " f"clip=[{cfg.sample_weight_min}, {cfg.sample_weight_max}]" ) print(f" curve_layout={curve_layout}") print(" note: 当前重点训练 pressure + derivative；可显式关闭 schedule 分支做固定制度对照") for epoch in range(1, cfg.epochs + 1): model.train() total = { "loss": 0.0, "loss_pressure": 0.0, "loss_derivative": 0.0, "loss_slope": 0.0, "loss_bias_pressure": 0.0, "loss_bias_derivative": 0.0, "loss_derivative_shape": 0.0, "loss_autofit_pressure": 0.0, "loss_autofit_derivative": 0.0, "sample_weight_mean": 0.0, "sample_weight_max": 0.0, } total_n = 0 for params_x, schedule_x, curve_y in train_loader: params_x = params_x.to(cfg.device) schedule_x = schedule_x.to(cfg.device) curve_y = curve_y.to(cfg.device) optimizer.zero_grad() pred = model_forward(model, params_x, schedule_x, cfg.use_schedule) losses = compute_weighted_loss( pred=pred, target=curve_y, slices=part_slices, curve_mean_raw=cfg.curve_mean_raw, curve_scale_raw=cfg.curve_scale_raw, huber_beta=cfg.huber_beta, w_pressure=cfg.w_pressure, w_derivative=cfg.w_derivative, w_slope=cfg.w_slope, w_bias_pressure=cfg.w_bias_pressure, w_bias_derivative=cfg.w_bias_derivative, w_derivative_shape=cfg.w_derivative_shape, w_autofit_pressure=cfg.w_autofit_pressure, w_autofit_derivative=cfg.w_autofit_derivative, use_sample_reweight=cfg.use_sample_reweight, sample_reweight_alpha=cfg.sample_reweight_alpha, sample_weight_min=cfg.sample_weight_min, sample_weight_max=cfg.sample_weight_max, ) losses["loss"].backward() optimizer.step() bs = params_x.size(0) for k in total: total[k] += losses[k].item() * bs total_n += bs denom = max(total_n, 1) train_metrics = {k: v / denom for k, v in total.items()} val_metrics = evaluate( model=model, loader=val_loader, device=cfg.device, slices=part_slices, cfg=cfg, ) row = {"epoch": epoch} for k, v in train_metrics.items(): row[f"train_{k}"] = float(v) for k, v in val_metrics.items(): row[f"val_{k}"] = float(v) history.append(row) print( f"[Epoch {epoch:03d}] " f"train={train_metrics['loss']:.6f} " f"(p={train_metrics['loss_pressure']:.6f}, " f"d={train_metrics['loss_derivative']:.6f}, " f"s={train_metrics['loss_slope']:.6f}, " f"bp={train_metrics['loss_bias_pressure']:.6f}, " f"bd={train_metrics['loss_bias_derivative']:.6f}, " f"ds={train_metrics['loss_derivative_shape']:.6f}, " f"ap={train_metrics['loss_autofit_pressure']:.6f}, " f"ad={train_metrics['loss_autofit_derivative']:.6f}, " f"wmean={train_metrics['sample_weight_mean']:.4f}, " f"wmax={train_metrics['sample_weight_max']:.4f}) " f"val={val_metrics['loss']:.6f} " f"(p={val_metrics['loss_pressure']:.6f}, " f"d={val_metrics['loss_derivative']:.6f}, " f"s={val_metrics['loss_slope']:.6f}, " f"bp={val_metrics['loss_bias_pressure']:.6f}, " f"bd={val_metrics['loss_bias_derivative']:.6f}, " f"ds={val_metrics['loss_derivative_shape']:.6f}, " f"ap={val_metrics['loss_autofit_pressure']:.6f}, " f"ad={val_metrics['loss_autofit_derivative']:.6f}, " f"wmean={val_metrics['sample_weight_mean']:.4f}, " f"wmax={val_metrics['sample_weight_max']:.4f})" ) if val_metrics["loss"] < best_val: best_val = val_metrics["loss"] torch.save( { "model_state_dict": model.state_dict(), "param_dim": param_dim, "schedule_dim": schedule_dim, "curve_dim": curve_dim, "hidden_dim": cfg.hidden_dim, "dropout": cfg.dropout, "use_schedule": cfg.use_schedule, "seed": int(cfg.seed), "curve_layout": curve_layout, "loss_weights": { "pressure": cfg.w_pressure, "derivative": cfg.w_derivative, "slope": cfg.w_slope, "bias_pressure": cfg.w_bias_pressure, "bias_derivative": cfg.w_bias_derivative, "derivative_shape": cfg.w_derivative_shape, "autofit_pressure": cfg.w_autofit_pressure, "autofit_derivative": cfg.w_autofit_derivative, }, "sample_reweight": { "enabled": cfg.use_sample_reweight, "alpha": cfg.sample_reweight_alpha, "weight_min": cfg.sample_weight_min, "weight_max": cfg.sample_weight_max, }, }, best_path, ) print(f" -> best model saved to: {best_path}") with open(cfg.output_dir / "history.json", "w", encoding="utf-8") as f: json.dump(history, f, ensure_ascii=False, indent=2) checkpoint = torch.load(best_path, map_location=cfg.device) best_model = ForwardSurrogate( param_dim=checkpoint["param_dim"], schedule_dim=checkpoint["schedule_dim"], curve_dim=checkpoint["curve_dim"], hidden_dim=checkpoint["hidden_dim"], dropout=checkpoint["dropout"], use_schedule=checkpoint.get("use_schedule", True), ).to(cfg.device) best_model.load_state_dict(checkpoint["model_state_dict"]) test_metrics = evaluate( model=best_model, loader=test_loader, device=cfg.device, slices=part_slices, cfg=cfg, ) print( f"[Final] test={test_metrics['loss']:.6f} " f"(p={test_metrics['loss_pressure']:.6f}, " f"d={test_metrics['loss_derivative']:.6f}, " f"s={test_metrics['loss_slope']:.6f}, " f"bp={test_metrics['loss_bias_pressure']:.6f}, " f"bd={test_metrics['loss_bias_derivative']:.6f}, " f"ds={test_metrics['loss_derivative_shape']:.6f}, " f"ap={test_metrics['loss_autofit_pressure']:.6f}, " f"ad={test_metrics['loss_autofit_derivative']:.6f}, " f"wmean={test_metrics['sample_weight_mean']:.4f}, " f"wmax={test_metrics['sample_weight_max']:.4f})" ) with open(cfg.output_dir / "metrics.json", "w", encoding="utf-8") as f: json.dump( { "best_val_loss": float(best_val), "test_metrics": {k: float(v) for k, v in test_metrics.items()}, "use_schedule": cfg.use_schedule, "seed": int(cfg.seed), "loss_weights": { "pressure": cfg.w_pressure, "derivative": cfg.w_derivative, "slope": cfg.w_slope, "bias_pressure": cfg.w_bias_pressure, "bias_derivative": cfg.w_bias_derivative, "derivative_shape": cfg.w_derivative_shape, "autofit_pressure": cfg.w_autofit_pressure, "autofit_derivative": cfg.w_autofit_derivative, }, "sample_reweight": { "enabled": cfg.use_sample_reweight, "alpha": cfg.sample_reweight_alpha, "weight_min": cfg.sample_weight_min, "weight_max": cfg.sample_weight_max, }, "curve_layout": curve_layout, }, f, ensure_ascii=False, indent=2, )