# -*- coding: utf-8 -*- """ eval_model.py - 加载训练好的 Transformer 模型（best_model.pt）和标准化参数（scaler.npz） - 对单个 CSV 进行标准化、滑窗切片并评测：wRMSE / MSE / MAE(original) - 可选输出最后一个窗口的预测与真值（反标准化）到 CSV 使用示例： uv run chaos_pdl/eval_model.py \ --csv_path chaos_pdl/data/achieve2/raw/IMG_1119_out_metrics.csv \ --model_path outputs/chaos/best_model.pt \ --scaler_path outputs/chaos/scaler.npz \ --seq_len 180 --stride 1 \ --save_dir outputs/chaos """ import argparse import os from typing import List import numpy as np import pandas as pd import torch # 复用训练脚本中的模块与评测函数 # 兼容两种导入方式： try: # 优先按包路径导入（若将项目安装为包或以工作区根为 sys.path） from chaos_pdl.pendulum_transformer import ( TimeSeriesTransformer, WeightedRMSELoss, WindowedTSDataset, ) except ModuleNotFoundError: # 退化为同目录导入（直接运行本脚本时更稳） from pendulum_transformer import ( TimeSeriesTransformer, WeightedRMSELoss, WindowedTSDataset, ) def build_windows_indices(T: int, seq_len: int, horizon: int, stride: int) -> np.ndarray: """与训练一致：生成起点索引，确保 s+seq_len+horizon 不越界。""" max_start = T - seq_len - horizon if max_start < 0: return np.array([], dtype=np.int64) return np.arange(0, max_start + 1, stride, dtype=np.int64) @torch.no_grad() def evaluate_losses(model, loader, device, weighted_rmse: WeightedRMSELoss): model.eval() total_wrmse, total_mse = 0.0, 0.0 count = 0 mse_fn = torch.nn.MSELoss() for xb, yb in loader: xb = xb.to(device) yb = yb.to(device) pred = model(xb) total_wrmse += weighted_rmse(pred, yb).item() * xb.size(0) total_mse += mse_fn(pred, yb).item() * xb.size(0) count += xb.size(0) return total_wrmse / max(1, count), total_mse / max(1, count) @torch.no_grad() def mae_in_original_scale(model, loader, device, mean: np.ndarray, std: np.ndarray): model.eval() total_mae, count = 0.0, 0 mean_t = torch.from_numpy(mean).to(device).view(1, 1, -1) std_t = torch.from_numpy(std).to(device).view(1, 1, -1) for xb, yb in loader: xb = xb.to(device) yb = yb.to(device) pred = model(xb) pred_denorm = pred * std_t + mean_t yb_denorm = yb * std_t + mean_t mae = torch.abs(pred_denorm - yb_denorm).mean() total_mae += mae.item() * xb.size(0) count += xb.size(0) return total_mae / max(1, count) def read_single_csv_align(path: str, expect_columns: List[str]) -> pd.DataFrame: """ 读取单个 CSV： - 若含 frame_index，则按其排序并移除； - 仅保留数值列； - 校验列集合与 scaler 中记录的列一致，并按预期顺序重排。 """ df = pd.read_csv(path) if "frame_index" in df.columns: df = df.sort_values("frame_index").drop(columns=["frame_index"]) # 对齐训练逻辑 df_num = df.select_dtypes(include=[np.number]).copy() has = list(df_num.columns) if set(has) != set(expect_columns): raise ValueError( f"CSV 数值列与训练时不一致\n 期望: {expect_columns}\n 实际: {has}\n" "请确保列集合一致（名称与类型），且来源一致。" ) # 按 scaler 的列顺序重排 df_num = df_num[expect_columns] return df_num.reset_index(drop=True) def main(): ap = argparse.ArgumentParser() ap.add_argument("--csv_path", type=str, required=True, help="待评测的 CSV 文件路径") ap.add_argument("--model_path", type=str, default="./outputs/chaos/best_model.pt") ap.add_argument("--scaler_path", type=str, default="./outputs/chaos/scaler.npz") ap.add_argument("--seq_len", type=int, default=180, help="历史窗口长度 L（需与训练一致）") ap.add_argument("--stride", type=int, default=1, help="评测滑窗步长") ap.add_argument("--batch_size", type=int, default=1024) ap.add_argument("--save_dir", type=str, default="", help="若提供，则保存最后一个窗口的预测/真值 CSV") ap.add_argument( "--plot_path", type=str, default="", help="每步RMSE(原尺度)曲线图的保存路径（.png）。若未提供且有save_dir，则保存到 save_dir/per_step_rmse.png", ) ap.add_argument( "--plot_heatmap", action="store_true", help="输出每步×每特征的RMSE热力图（原尺度），保存到 save_dir/per_step_feature_rmse.png 与 CSV", ) ap.add_argument( "--plot_feature", type=str, default="", help="指定一个特征名，对最后一个窗口画出预测与真值曲线；默认取第一个列名", ) args = ap.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") # 读取 scaler（包含列顺序） sc = np.load(args.scaler_path, allow_pickle=True) mean = sc["mean"] std = sc["std"] columns = list(sc["columns"].tolist()) print(f"Scaler loaded. Columns={columns}") # 读取并对齐 CSV df = read_single_csv_align(args.csv_path, columns) values = df.values.astype(np.float32) T, F = values.shape print(f"Data loaded. T={T}, F={F}") # 加载模型 ckpt = torch.load(args.model_path, map_location=device) cfg = ckpt["config"] loss_cfg = ckpt.get("loss_cfg", {"mode": "exp", "gamma": 0.99, "alpha": 0.02, "min_w": 0.05}) horizon = int(cfg["horizon"]) # 评测时以 ckpt 中的 horizon 为准 model = TimeSeriesTransformer( in_features=F, d_model=cfg["d_model"], nhead=cfg["nhead"], num_layers=cfg["num_layers"], dim_feedforward=cfg["dim_ff"], dropout=cfg["dropout"], horizon=horizon, ).to(device) model.load_state_dict(ckpt["model_state"]) model.eval() # 构造评测窗口 starts = build_windows_indices(T, args.seq_len, horizon, args.stride) if len(starts) == 0: raise ValueError( f"数据不足：T={T}，需要至少 seq_len({args.seq_len})+horizon({horizon}) 帧。" ) print(f"Eval windows: {len(starts)} (L={args.seq_len}, H={horizon}, stride={args.stride})") # 标准化（使用训练好的 mean/std） values_norm = (values - mean) / std # DataLoader ds = WindowedTSDataset(values_norm, starts, args.seq_len, horizon) loader = torch.utils.data.DataLoader(ds, batch_size=args.batch_size, shuffle=False) # 构造损失 weighted_rmse = WeightedRMSELoss( horizon=horizon, mode=loss_cfg.get("mode", "exp"), gamma=loss_cfg.get("gamma", 0.99), alpha=loss_cfg.get("alpha", 0.02), min_w=loss_cfg.get("min_w", 0.05), ).to(device) # 评测 val_wrmse, val_mse = evaluate_losses(model, loader, device, weighted_rmse) val_mae_orig = mae_in_original_scale(model, loader, device, mean, std) print( f"[Eval] wRMSE={val_wrmse:.6f} | MSE={val_mse:.6f} | MAE(original)={val_mae_orig:.6f}" ) # 统计每个预测步 t 的原尺度 RMSE，并可视化；可选统计每步×每特征的RMSE热力图 try: import matplotlib.pyplot as plt model.eval() H = horizon # 累积每步的 MSE（原尺度） step_sse = np.zeros(H, dtype=np.float64) step_count = 0 # 累积的样本数（B*F）按 batch 聚合 # 若需要热力图：累积 [H,F] 的平方误差和，以及样本计数（按窗口计数） per_feat_sse = None # shape [H,F] window_count = 0 mean_t = torch.from_numpy(mean).to(device).view(1, 1, -1) std_t = torch.from_numpy(std).to(device).view(1, 1, -1) with torch.no_grad(): for xb, yb in loader: xb = xb.to(device) yb = yb.to(device) pred = model(xb) pred_den = pred * std_t + mean_t yb_den = yb * std_t + mean_t err = pred_den - yb_den # [B,H,F] # 对 batch 和 feature 求均方，先累加平方和，最后再除以总数 sse_h = (err ** 2).sum(dim=(0, 2)).detach().cpu().numpy() # [H] step_sse += sse_h step_count += (xb.size(0) * xb.size(2)) if args.plot_heatmap: sse_hf = (err ** 2).sum(dim=0).detach().cpu().numpy() # [H,F] 按 batch 合并 if per_feat_sse is None: per_feat_sse = sse_hf.astype(np.float64) else: per_feat_sse += sse_hf window_count += xb.size(0) step_mse = step_sse / max(1, step_count) step_rmse = np.sqrt(np.maximum(step_mse, 1e-12)) # 保存数据到 CSV（若 save_dir 提供） if args.save_dir: os.makedirs(args.save_dir, exist_ok=True) csv_path = os.path.join(args.save_dir, "per_step_rmse.csv") pd.DataFrame({ "step": np.arange(1, H + 1), "rmse": step_rmse, "mse": step_mse, }).to_csv(csv_path, index=False) print(f"Saved: {csv_path}") # 绘图路径 plot_path = args.plot_path if not plot_path and args.save_dir: plot_path = os.path.join(args.save_dir, "per_step_rmse.png") if plot_path: plt.figure(figsize=(8, 4.5)) plt.plot(np.arange(1, H + 1), step_rmse, marker="o", lw=1.5) plt.xlabel("Prediction step (t)") plt.ylabel("RMSE (original scale)") plt.title("Per-step RMSE vs prediction step") plt.grid(True, alpha=0.3) plt.tight_layout() plt.savefig(plot_path, dpi=150) print(f"Saved: {plot_path}") # 输出热力图与CSV if args.plot_heatmap and per_feat_sse is not None and args.save_dir: os.makedirs(args.save_dir, exist_ok=True) per_feat_mse = per_feat_sse / max(1, window_count) # [H,F] 每步每特征 MSE（按窗口平均） per_feat_rmse = np.sqrt(np.maximum(per_feat_mse, 1e-12)) # CSV：行为 step，列为各特征 csv_hm = os.path.join(args.save_dir, "per_step_feature_rmse.csv") pd.DataFrame(per_feat_rmse, columns=columns, index=np.arange(1, H + 1)).to_csv(csv_hm) print(f"Saved: {csv_hm}") # 热力图 plt.figure(figsize=(min(12, 1.5 + 0.35 * len(columns) + 6), 5)) im = plt.imshow(per_feat_rmse, aspect="auto", origin="lower", cmap="viridis") plt.colorbar(im, label="RMSE") plt.yticks(ticks=np.arange(0, H, max(1, H // 10)), labels=(np.arange(1, H + 1)[::max(1, H // 10)])) plt.xticks(ticks=np.arange(len(columns)), labels=columns, rotation=45, ha="right") plt.xlabel("Feature") plt.ylabel("Prediction step (t)") plt.title("Per-step per-feature RMSE (original scale)") plt.tight_layout() hm_path = os.path.join(args.save_dir, "per_step_feature_rmse.png") plt.savefig(hm_path, dpi=150) print(f"Saved: {hm_path}") except Exception as e: print(f"[warn] 每步RMSE绘图失败：{e}") # 可选：保存最后一个窗口的预测与真值（反标准化） if args.save_dir: os.makedirs(args.save_dir, exist_ok=True) last_start = int(starts[-1]) x_last = values_norm[last_start : last_start + args.seq_len] y_true = values[last_start + args.seq_len : last_start + args.seq_len + horizon] with torch.no_grad(): pred = ( model(torch.from_numpy(x_last).unsqueeze(0).to(device)) .cpu() .numpy()[0] ) pred_denorm = pred * std + mean # 保存为两份 CSV：pred 和 true，行是步长 t=1..H，列为各特征名 pred_df = pd.DataFrame(pred_denorm, columns=columns) true_df = pd.DataFrame(y_true, columns=columns) pred_path = os.path.join(args.save_dir, "last_window_pred.csv") true_path = os.path.join(args.save_dir, "last_window_true.csv") pred_df.to_csv(pred_path, index=False) true_df.to_csv(true_path, index=False) print(f"Saved: {pred_path}\nSaved: {true_path}") # 可选：对某个特征画最后一个窗口的预测-真值曲线 try: import matplotlib.pyplot as plt feat = args.plot_feature or columns[0] if feat not in columns: print(f"[warn] plot_feature='{feat}' 不在列中，使用默认 {columns[0]}") feat = columns[0] fi = columns.index(feat) plt.figure(figsize=(7, 4)) plt.plot(np.arange(1, horizon + 1), pred_denorm[:, fi], label="pred", marker="o", lw=1.5) plt.plot(np.arange(1, horizon + 1), y_true[:, fi], label="true", marker="o", lw=1.5) plt.xlabel("Prediction step (t)") plt.ylabel(feat) plt.title(f"Last window prediction vs truth: {feat}") plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() line_path = os.path.join(args.save_dir, f"last_window_{feat}_pred_vs_true.png") plt.savefig(line_path, dpi=150) print(f"Saved: {line_path}") except Exception as e: print(f"[warn] 最后窗口曲线绘图失败：{e}") if __name__ == "__main__": main()