first commit
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
uv run pendulum_transformer.py --csv data/448.csv data/618.csv data/548.csv \
|
||||
--epochs 10 --batch 64 --eval_all\
|
||||
--s_in 120 --s_out 30 --stride 5 \
|
||||
--ckpt ckpt/pendulum_tf_base.pt
|
||||
+10811
File diff suppressed because it is too large
Load Diff
+22406
File diff suppressed because it is too large
Load Diff
+26624
File diff suppressed because it is too large
Load Diff
+34496
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,116 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
几何与坐标变换工具：
|
||||
- 基于最小二乘的圆拟合
|
||||
- 像素坐标 <-> 单位圆 (sin, cos) 变换
|
||||
|
||||
与训练脚本的定义保持一致：
|
||||
theta = atan2(x - cx, cy - y)
|
||||
sin = sin(theta) = (x - cx) / r
|
||||
cos = cos(theta) = (cy - y) / r = -(y - cy) / r
|
||||
|
||||
因此，反变换：
|
||||
x = cx + r * sin
|
||||
y = cy - r * cos
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import numpy as np
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
def fit_circle_lstsq(x: np.ndarray, y: np.ndarray) -> Tuple[float, float, float]:
|
||||
"""最小二乘拟合圆：解 x^2 + y^2 = 2 a x + 2 b y + c
|
||||
|
||||
返回 (cx, cy, r)
|
||||
"""
|
||||
x = np.asarray(x, dtype=np.float64)
|
||||
y = np.asarray(y, dtype=np.float64)
|
||||
A = np.c_[2 * x, 2 * y, np.ones_like(x)]
|
||||
b = x * x + y * y
|
||||
sol, *_ = np.linalg.lstsq(A, b, rcond=None)
|
||||
a, b_, c = sol
|
||||
cx, cy = a, b_
|
||||
r = math.sqrt(max(cx * cx + cy * cy + c, 0.0))
|
||||
return float(cx), float(cy), float(r)
|
||||
|
||||
|
||||
def _estimate_center_radius(seq_xy: np.ndarray) -> Tuple[float, float, float]:
|
||||
"""对单条序列估计 (cx, cy, r)。
|
||||
- 当样本点 >=3 时使用最小二乘拟合
|
||||
- 否则使用均值点作为中心、均距为半径；若半径过小，则回退为 1.0
|
||||
"""
|
||||
seq_xy = np.asarray(seq_xy, dtype=np.float64)
|
||||
S = seq_xy.shape[0]
|
||||
if S >= 3:
|
||||
cx, cy, r = fit_circle_lstsq(seq_xy[:, 0], seq_xy[:, 1])
|
||||
else:
|
||||
cx, cy = np.mean(seq_xy, axis=0)
|
||||
d = np.linalg.norm(seq_xy - np.array([cx, cy], dtype=np.float64), axis=1)
|
||||
r = float(np.median(d) if d.size else 1.0)
|
||||
r = float(max(r, 1e-6))
|
||||
return float(cx), float(cy), r
|
||||
|
||||
|
||||
def pixels_to_unit_sincos(seq_xy: np.ndarray) -> Tuple[np.ndarray, Tuple[float, float, float]]:
|
||||
"""将单条像素轨迹 [S,2] 转为单位 (sin,cos) [S,2]，并返回 (cx,cy,r)。"""
|
||||
seq_xy = np.asarray(seq_xy, dtype=np.float64)
|
||||
cx, cy, r = _estimate_center_radius(seq_xy)
|
||||
x = seq_xy[:, 0]
|
||||
y = seq_xy[:, 1]
|
||||
# 与训练一致的角度定义
|
||||
theta = np.arctan2(x - cx, cy - y)
|
||||
sc = np.stack([np.sin(theta), np.cos(theta)], axis=-1).astype(np.float32)
|
||||
return sc, (float(cx), float(cy), float(r))
|
||||
|
||||
|
||||
def batch_pixels_to_unit_sincos(batch_xy: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""批量像素 -> 单位 (sin,cos)。
|
||||
|
||||
参数：
|
||||
batch_xy: [B,S,2] 像素坐标
|
||||
返回：
|
||||
sc: [B,S,2] 单位 (sin,cos)
|
||||
centers: [B,2] (cx,cy)
|
||||
radii: [B] r
|
||||
"""
|
||||
batch_xy = np.asarray(batch_xy, dtype=np.float64)
|
||||
B, S, _ = batch_xy.shape
|
||||
sc = np.empty((B, S, 2), dtype=np.float32)
|
||||
centers = np.empty((B, 2), dtype=np.float32)
|
||||
radii = np.empty((B,), dtype=np.float32)
|
||||
for i in range(B):
|
||||
seq = batch_xy[i]
|
||||
sc_i, (cx, cy, r) = pixels_to_unit_sincos(seq)
|
||||
sc[i] = sc_i
|
||||
centers[i] = (cx, cy)
|
||||
radii[i] = r
|
||||
return sc, centers, radii
|
||||
|
||||
|
||||
def unit_sincos_to_pixels(sc: np.ndarray, centers: np.ndarray, radii: np.ndarray) -> np.ndarray:
|
||||
"""单位 (sin,cos) -> 像素坐标。
|
||||
|
||||
支持广播：
|
||||
sc: [B,S,2]
|
||||
centers: [B,2]
|
||||
radii: [B] or [B,1]
|
||||
返回：
|
||||
pix: [B,S,2]
|
||||
"""
|
||||
sc = np.asarray(sc, dtype=np.float64)
|
||||
centers = np.asarray(centers, dtype=np.float64)
|
||||
radii = np.asarray(radii, dtype=np.float64).reshape(-1, 1, 1) # [B,1,1]
|
||||
|
||||
sin = sc[..., 0]
|
||||
cos = sc[..., 1]
|
||||
cx = centers[:, 0].reshape(-1, 1)
|
||||
cy = centers[:, 1].reshape(-1, 1)
|
||||
|
||||
x = cx + radii[:, 0, 0:1] * sin # 广播 [B,1] * [B,S]
|
||||
y = cy - radii[:, 0, 0:1] * cos
|
||||
|
||||
pix = np.stack([x, y], axis=-1)
|
||||
return pix.astype(np.float32)
|
||||
@@ -0,0 +1,479 @@
|
||||
# pendulum_transformer.py
|
||||
import argparse, math, os
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
# ============ 工具函数 ============
|
||||
def moving_average(a, w=5):
|
||||
if w <= 1: return a
|
||||
pad = np.pad(a, (w//2, w-1-w//2), mode='edge')
|
||||
k = np.ones(w)/w
|
||||
return np.convolve(pad, k, mode='valid')
|
||||
|
||||
def fit_circle_lstsq(x, y):
|
||||
# 解: x^2 + y^2 = 2 a x + 2 b y + c
|
||||
A = np.c_[2*x, 2*y, np.ones_like(x)]
|
||||
b = x**2 + y**2
|
||||
sol, *_ = np.linalg.lstsq(A, b, rcond=None)
|
||||
a, b_, c = sol
|
||||
cx, cy = a, b_
|
||||
r = math.sqrt(max(cx*cx + cy*cy + c, 0.0))
|
||||
return cx, cy, r
|
||||
|
||||
def wrap_to_pi(a):
|
||||
return (a + np.pi) % (2*np.pi) - np.pi
|
||||
|
||||
def build_windows(arr, s_in=120, s_out=30, stride=10):
|
||||
# arr: [T, F]
|
||||
T, F = arr.shape
|
||||
X, Y, t_idx = [], [], []
|
||||
for t in range(0, T - s_in - s_out + 1, stride):
|
||||
X.append(arr[t:t+s_in])
|
||||
Y.append(arr[t+s_in:t+s_in+s_out])
|
||||
t_idx.append(t)
|
||||
return np.array(X), np.array(Y), np.array(t_idx) # [N, S, F], [N, S, F], [N]
|
||||
|
||||
def sincos_to_angle(s, c):
|
||||
return np.arctan2(s, c)
|
||||
|
||||
def angle_rmse(pred_sc, true_sc):
|
||||
# pred_sc, true_sc: [..., 2] -> (sin, cos)
|
||||
ps, pc = pred_sc[...,0], pred_sc[...,1]
|
||||
ts, tc = true_sc[...,0], true_sc[...,1]
|
||||
pa = np.arctan2(ps, pc)
|
||||
ta = np.arctan2(ts, tc)
|
||||
diff = wrap_to_pi(pa - ta)
|
||||
return np.sqrt(np.mean(diff**2))
|
||||
|
||||
def split_even_indices(n, ratios=(0.7, 0.15, 0.15), block=50, seed=123):
|
||||
"""将 [0..n-1] 均匀地分成三部分（train/val/test）"""
|
||||
if n <= 0:
|
||||
return np.array([], dtype=int), np.array([], dtype=int), np.array([], dtype=int)
|
||||
probs = np.array(ratios, dtype=np.float64)
|
||||
s = probs.sum()
|
||||
probs = probs / (s if s > 0 else 1.0)
|
||||
rng = np.random.default_rng(seed)
|
||||
idxs = np.arange(n)
|
||||
tr, va, te = [], [], []
|
||||
for start in range(0, n, block):
|
||||
end = min(start + block, n)
|
||||
m = end - start
|
||||
target = probs * m
|
||||
base = np.floor(target).astype(int)
|
||||
rem = m - int(base.sum())
|
||||
if rem > 0:
|
||||
frac = target - base
|
||||
order = np.argsort(-(frac + 1e-9 * rng.random(frac.shape)))
|
||||
for k in order[:rem]:
|
||||
base[k] += 1
|
||||
local = idxs[start:end].copy()
|
||||
rng.shuffle(local)
|
||||
a, b, c = int(base[0]), int(base[1]), int(base[2])
|
||||
tr.extend(local[:a]); va.extend(local[a:a+b]); te.extend(local[a+b:a+b+c])
|
||||
return np.array(tr, dtype=int), np.array(va, dtype=int), np.array(te, dtype=int)
|
||||
|
||||
# ============ 数据集 ============
|
||||
class SeqDataset(Dataset):
|
||||
def __init__(self, X, Y, mu=None, std=None):
|
||||
self.X = X.astype(np.float32)
|
||||
self.Y = Y.astype(np.float32)
|
||||
if mu is None or std is None:
|
||||
mu = self.X.reshape(-1, self.X.shape[-1]).mean(0, keepdims=True)
|
||||
std = self.X.reshape(-1, self.X.shape[-1]).std(0, keepdims=True) + 1e-8
|
||||
self.mu = mu.astype(np.float32)
|
||||
self.std = std.astype(np.float32)
|
||||
|
||||
def __len__(self): return self.X.shape[0]
|
||||
def __getitem__(self, i):
|
||||
x = (self.X[i] - self.mu) / self.std
|
||||
y = (self.Y[i] - self.mu) / self.std
|
||||
return torch.from_numpy(x), torch.from_numpy(y)
|
||||
|
||||
# ============ 位置编码 ============
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model, max_len=4096):
|
||||
super().__init__()
|
||||
pe = torch.zeros(max_len, d_model)
|
||||
pos = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
|
||||
div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model))
|
||||
pe[:, 0::2] = torch.sin(pos * div)
|
||||
pe[:, 1::2] = torch.cos(pos * div)
|
||||
self.register_buffer('pe', pe.unsqueeze(0)) # [1, L, D]
|
||||
def forward(self, x):
|
||||
return x + self.pe[:, :x.size(1), :]
|
||||
|
||||
# ============ 极简 Transformer ============
|
||||
class TimeSeriesTransformer(nn.Module):
|
||||
def __init__(self, in_dim=2, d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512, dropout=0.1, out_dim=2):
|
||||
super().__init__()
|
||||
self.in_proj = nn.Linear(in_dim, d_model)
|
||||
self.out_proj = nn.Linear(d_model, out_dim)
|
||||
self.tgt_proj = nn.Linear(out_dim, d_model)
|
||||
self.pos_enc = PositionalEncoding(d_model)
|
||||
self.tf = nn.Transformer(
|
||||
d_model=d_model, nhead=nhead,
|
||||
num_encoder_layers=enc_layers, num_decoder_layers=dec_layers,
|
||||
dim_feedforward=d_ff, dropout=dropout, batch_first=True
|
||||
)
|
||||
|
||||
def forward(self, src, tgt_in, tgt_mask=None):
|
||||
src = self.pos_enc(self.in_proj(src))
|
||||
tgt = self.pos_enc(self.tgt_proj(tgt_in))
|
||||
out = self.tf(src, tgt, tgt_mask=tgt_mask) # [B, S_out, D]
|
||||
return self.out_proj(out)
|
||||
|
||||
# ============ 主流程 ============
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--csv", type=str, nargs='+', required=True, help="CSV files")
|
||||
ap.add_argument("--split", type=float, nargs=3, default=[0.7, 0.15, 0.15])
|
||||
ap.add_argument("--s_in", type=int, default=180)
|
||||
ap.add_argument("--s_out", type=int, default=30)
|
||||
ap.add_argument("--stride", type=int, default=10)
|
||||
ap.add_argument("--epochs", type=int, default=10)
|
||||
ap.add_argument("--batch", type=int, default=64)
|
||||
ap.add_argument("--lr", type=float, default=1e-3)
|
||||
ap.add_argument("--ckpt", type=str, default="ckpt/pendulum_tf_base.pt")
|
||||
ap.add_argument("--eval_only", action="store_true")
|
||||
# 新增：评估相关
|
||||
ap.add_argument("--eval_all", action="store_true", help="Evaluate on the entire test set")
|
||||
ap.add_argument("--eval_samples", type=int, default=6, help="How many random samples to plot from test set")
|
||||
ap.add_argument("--seed", type=int, default=123, help="Random seed for eval sampling")
|
||||
args = ap.parse_args()
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# 规范化路径
|
||||
def _norm(p: str) -> str:
|
||||
return os.path.abspath(os.path.expanduser(p))
|
||||
|
||||
csv_list = sorted(set(_norm(p) for p in (args.csv or [])))
|
||||
if not csv_list:
|
||||
raise ValueError("No CSV files provided.")
|
||||
|
||||
ratios = np.array(args.split, dtype=np.float64)
|
||||
s = ratios.sum()
|
||||
ratios = ratios / (s if s > 0 else 1.0)
|
||||
print(f"[Split] auto per-CSV even mode: files={len(csv_list)} | ratios train/val/test = {ratios.tolist()}")
|
||||
|
||||
Xtr_list, Ytr_list, Xva_list, Yva_list, Xte_list, Yte_list = [], [], [], [], [], []
|
||||
|
||||
for i, csv_path in enumerate(csv_list):
|
||||
df = pd.read_csv(csv_path)
|
||||
x = df['x'].to_numpy(dtype=np.float64)
|
||||
y = df['y'].to_numpy(dtype=np.float64)
|
||||
|
||||
# 圆拟合 => 悬点
|
||||
cx, cy, r = fit_circle_lstsq(x, y)
|
||||
print(f"[{i+1}/{len(csv_list)}] {os.path.basename(csv_path)} pivot=({cx:.2f},{cy:.2f}), r~{r:.2f}")
|
||||
|
||||
# 像素 -> 角度；竖直向下为 0
|
||||
theta = np.arctan2(x - cx, cy - y)
|
||||
theta = np.unwrap(theta)
|
||||
sin_t, cos_t = np.sin(theta), np.cos(theta)
|
||||
feats = np.stack([sin_t, cos_t], axis=1) # [T, 2]
|
||||
|
||||
# 滑窗
|
||||
Xi, Yi, _ = build_windows(feats, s_in=args.s_in, s_out=args.s_out, stride=args.stride)
|
||||
if Xi.size == 0:
|
||||
continue
|
||||
tr_idx, va_idx, te_idx = split_even_indices(Xi.shape[0], ratios)
|
||||
if tr_idx.size: Xtr_list.append(Xi[tr_idx]); Ytr_list.append(Yi[tr_idx])
|
||||
if va_idx.size: Xva_list.append(Xi[va_idx]); Yva_list.append(Yi[va_idx])
|
||||
if te_idx.size: Xte_list.append(Xi[te_idx]); Yte_list.append(Yi[te_idx])
|
||||
|
||||
def _concat(lst, steps):
|
||||
if not lst:
|
||||
return np.zeros((0, steps, 2), dtype=np.float32)
|
||||
return np.concatenate(lst, axis=0).astype(np.float32)
|
||||
|
||||
Xtr = _concat(Xtr_list, args.s_in)
|
||||
Ytr = _concat(Ytr_list, args.s_out)
|
||||
Xva = _concat(Xva_list, args.s_in)
|
||||
Yva = _concat(Yva_list, args.s_out)
|
||||
Xte = _concat(Xte_list, args.s_in)
|
||||
Yte = _concat(Yte_list, args.s_out)
|
||||
|
||||
if Xtr.shape[0] == 0: raise ValueError("No training samples produced.")
|
||||
if Xva.shape[0] == 0: raise ValueError("No validation samples produced.")
|
||||
if Xte.shape[0] == 0: raise ValueError("No test samples produced.")
|
||||
|
||||
# 标准化参数（仅基于训练集）
|
||||
base = Xtr
|
||||
mu = base.reshape(-1, base.shape[-1]).mean(0, keepdims=True).astype(np.float32)
|
||||
std = base.reshape(-1, base.shape[-1]).std(0, keepdims=True).astype(np.float32) + 1e-8
|
||||
|
||||
dtr = SeqDataset(Xtr, Ytr, mu, std)
|
||||
dva = SeqDataset(Xva, Yva, mu, std)
|
||||
dte = SeqDataset(Xte, Yte, mu, std)
|
||||
|
||||
nw = 2
|
||||
Ltr = DataLoader(dtr, batch_size=args.batch, shuffle=True, drop_last=True, pin_memory=True, num_workers=nw)
|
||||
Lva = DataLoader(dva, batch_size=args.batch, shuffle=False, pin_memory=True, num_workers=nw)
|
||||
Lte = DataLoader(dte, batch_size=args.batch, shuffle=False, pin_memory=True, num_workers=nw)
|
||||
|
||||
cfg = dict(d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512)
|
||||
net = TimeSeriesTransformer(in_dim=2, out_dim=2, **cfg).to(device)
|
||||
|
||||
mse = nn.MSELoss()
|
||||
opt = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)
|
||||
|
||||
# ---- warmup + cosine 调度（按 step）----
|
||||
total_steps = max(1, args.epochs * len(Ltr))
|
||||
warmup = min(500, max(1, total_steps // 10))
|
||||
def lr_lambda(step):
|
||||
if step < warmup:
|
||||
return float(step + 1) / float(warmup)
|
||||
progress = (step - warmup) / max(1, total_steps - warmup)
|
||||
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
|
||||
min_factor = 0.01
|
||||
return min_factor + (1 - min_factor) * cosine
|
||||
scheduler = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda=lr_lambda)
|
||||
|
||||
# === 标准化 <-> 原始空间 转换 & 单位化 ===
|
||||
mu_t = torch.from_numpy(mu).to(device) # [1, 2]
|
||||
std_t = torch.from_numpy(std).to(device) # [1, 2]
|
||||
|
||||
def to_raw(z): return z * std_t + mu_t
|
||||
def to_std(z): return (z - mu_t) / std_t
|
||||
def unitize(z, eps=1e-6):
|
||||
n = torch.clamp(torch.linalg.norm(z, dim=-1, keepdim=True), min=eps)
|
||||
return z / n
|
||||
|
||||
# === 预测一个 batch 的自回归函数（返回 raw 空间预测） ===
|
||||
def rollout_raw(xb_std, steps):
|
||||
# xb_std: [B, S_in, 2] 标准化空间
|
||||
tgt_tokens = xb_std[:, -1:, :2].clone() # 起始 token（标准化空间）
|
||||
for _ in range(steps):
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(tgt_tokens.size(1)).to(device)
|
||||
out = net(xb_std, tgt_tokens, tgt_mask=mask) # 标准化
|
||||
next_std = out[:, -1:, :]
|
||||
next_raw = to_raw(next_std) # -> 原始
|
||||
next_raw = unitize(next_raw) # 单位圆约束
|
||||
next_std = to_std(next_raw) # -> 标准化（与训练分布一致）
|
||||
tgt_tokens = torch.cat([tgt_tokens, next_std], dim=1)
|
||||
pred_std = tgt_tokens[:, 1:, :] # [B, steps, 2] 标准化
|
||||
return to_raw(pred_std) # -> 原始
|
||||
|
||||
# === 改动一：闭环 rollout 损失所需的角度/方向损失 ===
|
||||
def cosine_ang_loss(pred_sc_raw, true_sc_raw, eps=1e-6):
|
||||
"""
|
||||
pred_sc_raw/true_sc_raw: [B,S,2]，分别是 RAW 空间下的 (sin,cos)
|
||||
计算 1 - cos(Δθ) 的平均值，忽略向量长度，只关注方向（角度）差。
|
||||
"""
|
||||
pu = pred_sc_raw / torch.clamp(torch.linalg.norm(pred_sc_raw, dim=-1, keepdim=True), min=eps)
|
||||
tu = true_sc_raw / torch.clamp(torch.linalg.norm(true_sc_raw, dim=-1, keepdim=True), min=eps)
|
||||
return (1.0 - (pu * tu).sum(dim=-1)).mean()
|
||||
|
||||
# === 训练 or 加载 ===
|
||||
os.makedirs(os.path.dirname(args.ckpt) or ".", exist_ok=True)
|
||||
if args.eval_only:
|
||||
if not os.path.exists(args.ckpt):
|
||||
raise FileNotFoundError(f"Checkpoint not found: {args.ckpt}. 请先训练一次或指定正确路径。")
|
||||
# 关键修复：显式关闭 weights_only 的安全反序列化限制（你信任本地 ckpt）
|
||||
payload = torch.load(args.ckpt, map_location=device, weights_only=False)
|
||||
# 若 ckpt 中包含 cfg，用 ckpt 的配置重建网络，避免结构不一致
|
||||
ckpt_cfg = payload.get("cfg", None)
|
||||
if ckpt_cfg is not None:
|
||||
# 如果与当前构造的不同，则重建
|
||||
if ckpt_cfg != cfg:
|
||||
net = TimeSeriesTransformer(in_dim=2, out_dim=2, **ckpt_cfg).to(device)
|
||||
net.load_state_dict(payload["model"], strict=True)
|
||||
|
||||
mu = payload["mu"]; std = payload["std"]
|
||||
mu_t = torch.from_numpy(mu).to(device); std_t = torch.from_numpy(std).to(device)
|
||||
print(f"Loaded checkpoint from {args.ckpt}")
|
||||
else:
|
||||
# Sanity：原始空间下 ||yb|| 应 ~1
|
||||
with torch.no_grad():
|
||||
xb0, yb0 = next(iter(Ltr))
|
||||
m_norm_raw = torch.linalg.norm(to_raw(yb0.to(device)), dim=-1).mean().item()
|
||||
print(f"[Sanity] mean ||yb|| in RAW space ~= {m_norm_raw:.4f}")
|
||||
|
||||
best_val = float("inf")
|
||||
global_step = 0
|
||||
for ep in range(1, args.epochs+1):
|
||||
net.train()
|
||||
total = 0.0
|
||||
total_tf = 0.0
|
||||
total_roll = 0.0
|
||||
|
||||
# ρ：闭环损失权重（线性升到 1.0；如需更稳可改为升到 0.5）
|
||||
rho = min(1.0, ep / max(1, args.epochs // 2))
|
||||
|
||||
for xb, yb in Ltr:
|
||||
xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
|
||||
B, S, F = yb.shape
|
||||
|
||||
# ---------- (A) Teacher Forcing 分支（与原来一致） ----------
|
||||
x_last = xb[:, -1:, :F]
|
||||
tgt_in = torch.cat([x_last, yb[:, :-1, :]], dim=1)
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(S).to(device)
|
||||
|
||||
pred_tf = net(xb, tgt_in, tgt_mask=mask) # 标准化空间
|
||||
pred_tf_u = unitize(to_raw(pred_tf)) # 原始空间单位圆（只对方向约束）
|
||||
yb_u = unitize(to_raw(yb))
|
||||
loss_tf = mse(pred_tf_u, yb_u)
|
||||
|
||||
# ---------- (B) 闭环 rollout 分支（改动一的核心） ----------
|
||||
with torch.no_grad():
|
||||
y_true_raw = to_raw(yb) # RAW，用于对齐真值方向
|
||||
y_pred_raw = rollout_raw(xb, steps=S) # RAW，模型自回归滚动预测
|
||||
loss_roll = cosine_ang_loss(y_pred_raw, y_true_raw)
|
||||
|
||||
# ---------- 合成总损失 ----------
|
||||
loss = (1.0 - rho) * loss_tf + rho * loss_roll
|
||||
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)
|
||||
opt.step()
|
||||
scheduler.step()
|
||||
|
||||
total += loss.item() * xb.size(0)
|
||||
total_tf += loss_tf.item() * xb.size(0)
|
||||
total_roll += loss_roll.item() * xb.size(0)
|
||||
global_step += 1
|
||||
|
||||
tr_loss = total / len(dtr)
|
||||
tr_tf = total_tf / len(dtr)
|
||||
tr_rl = total_roll / len(dtr)
|
||||
|
||||
# 验证：保持原先的 TF-式度量（也可扩展加闭环验证，这里先不改动）
|
||||
net.eval()
|
||||
with torch.no_grad():
|
||||
total_v = 0.0
|
||||
for xb, yb in Lva:
|
||||
xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
|
||||
x_last = xb[:, -1:, :2]
|
||||
tgt_in = torch.cat([x_last, yb[:, :-1, :]], dim=1)
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(yb.size(1)).to(device)
|
||||
pred = net(xb, tgt_in, tgt_mask=mask)
|
||||
pred_u = unitize(to_raw(pred))
|
||||
yb_u = unitize(to_raw(yb))
|
||||
total_v += mse(pred_u, yb_u).item() * xb.size(0)
|
||||
va_loss = total_v / len(dva)
|
||||
|
||||
cur_lr = scheduler.get_last_lr()[0]
|
||||
print(f"[Epoch {ep:02d}] ρ={rho:.2f} | train TOTAL={tr_loss:.6f} (TF={tr_tf:.6f}, ROLL={tr_rl:.6f}) | val MSE(u_raw)={va_loss:.6f} | lr={cur_lr:.6e}")
|
||||
|
||||
if not np.isnan(va_loss) and va_loss < best_val:
|
||||
best_val = va_loss
|
||||
payload = {"model": net.state_dict(), "mu": mu, "std": std, "cfg": cfg}
|
||||
torch.save(payload, args.ckpt)
|
||||
print(f"Saved checkpoint -> {args.ckpt}")
|
||||
|
||||
# ===== 评估 =====
|
||||
net.eval()
|
||||
rng = np.random.default_rng(args.seed)
|
||||
|
||||
with torch.no_grad():
|
||||
# --- 1) 全量评估（可选） ---
|
||||
if args.eval_all:
|
||||
steps = args.s_out
|
||||
sum_sq = 0.0
|
||||
n_total = 0
|
||||
per_step_sq = np.zeros(steps, dtype=np.float64)
|
||||
total_samples = 0
|
||||
|
||||
for xb, yb in Lte:
|
||||
xb = xb.to(device, non_blocking=True); yb = yb.to(device, non_blocking=True)
|
||||
y_pred_raw = rollout_raw(xb, steps) # [B, steps, 2] raw
|
||||
y_true_raw = to_raw(yb) # [B, steps, 2] raw
|
||||
y_pred_np = y_pred_raw.cpu().numpy()
|
||||
y_true_np = y_true_raw.cpu().numpy()
|
||||
|
||||
# 单位圆 -> 角度误差
|
||||
n1 = np.clip(np.linalg.norm(y_pred_np, axis=-1, keepdims=True), 1e-6, None)
|
||||
n2 = np.clip(np.linalg.norm(y_true_np, axis=-1, keepdims=True), 1e-6, None)
|
||||
y_pred_u = y_pred_np / n1
|
||||
y_true_u = y_true_np / n2
|
||||
|
||||
pa = np.arctan2(y_pred_u[...,0], y_pred_u[...,1]) # 注意顺序 (sin,cos)
|
||||
ta = np.arctan2(y_true_u[...,0], y_true_u[...,1])
|
||||
diff = wrap_to_pi(pa - ta) # [B, steps]
|
||||
|
||||
sum_sq += np.sum(diff**2)
|
||||
per_step_sq += np.sum(diff**2, axis=0)
|
||||
n_total += diff.size
|
||||
total_samples += diff.shape[0]
|
||||
|
||||
overall_rmse = math.sqrt(sum_sq / n_total)
|
||||
per_step_rmse = np.sqrt(per_step_sq / max(1, total_samples))
|
||||
|
||||
print(f"[Test:ALL] overall angle RMSE = {overall_rmse:.6f} rad ({overall_rmse*180/np.pi:.3f} deg)")
|
||||
# 保存每步 RMSE 曲线
|
||||
ts = np.arange(1, steps+1)
|
||||
plt.figure(figsize=(6,4))
|
||||
plt.plot(ts, per_step_rmse * 180/np.pi, marker='o')
|
||||
plt.xlabel("horizon (step)"); plt.ylabel("RMSE (deg)")
|
||||
plt.title("Per-step angle RMSE on test set")
|
||||
plt.grid(True, linestyle="--", alpha=0.5)
|
||||
plt.tight_layout()
|
||||
plt.savefig("rmse_per_step.png", dpi=150)
|
||||
print("Saved plot -> rmse_per_step.png")
|
||||
|
||||
# --- 2) 随机抽样 K 个样本可视化 ---
|
||||
X_list, Y_list = [], []
|
||||
for xb, yb in Lte:
|
||||
X_list.append(xb); Y_list.append(yb)
|
||||
X_all = torch.cat(X_list, dim=0).to(device, non_blocking=True)
|
||||
Y_all = torch.cat(Y_list, dim=0).to(device, non_blocking=True)
|
||||
|
||||
N = X_all.size(0)
|
||||
K = min(max(0, args.eval_samples), N)
|
||||
if K > 0:
|
||||
idxs = rng.choice(N, size=K, replace=False)
|
||||
steps = Y_all.size(1)
|
||||
y_pred_raw_all = rollout_raw(X_all[idxs], steps) # [K, steps, 2]
|
||||
y_true_raw_all = to_raw(Y_all[idxs]) # [K, steps, 2]
|
||||
|
||||
def unit_np(z):
|
||||
n = np.linalg.norm(z, axis=-1, keepdims=True)
|
||||
n = np.clip(n, 1e-6, None)
|
||||
return z / n
|
||||
|
||||
y_pred_np = y_pred_raw_all.cpu().numpy()
|
||||
y_true_np = y_true_raw_all.cpu().numpy()
|
||||
y_pred_u = unit_np(y_pred_np)
|
||||
y_true_u = unit_np(y_true_np)
|
||||
|
||||
for j in range(K):
|
||||
ya = np.arctan2(y_true_u[j,:,0], y_true_u[j,:,1])
|
||||
pa = np.arctan2(y_pred_u[j,:,0], y_pred_u[j,:,1])
|
||||
pa_aligned = pa + 2*np.pi * np.round((ya - pa)/(2*np.pi))
|
||||
|
||||
plt.figure(figsize=(8,4))
|
||||
plt.plot(np.arange(steps), ya, label="true θ")
|
||||
plt.plot(np.arange(steps), pa_aligned, '--', label="pred θ (aligned)")
|
||||
plt.xlabel("future step"); plt.ylabel("angle (rad)")
|
||||
plt.title(f"{steps}-step forecast (test sample {j})")
|
||||
plt.legend(); plt.tight_layout()
|
||||
out_png = f"pred_plot_sample_{j}.png"
|
||||
plt.savefig(out_png, dpi=150)
|
||||
plt.close()
|
||||
print(f"Saved plot -> {out_png}")
|
||||
|
||||
# 如果没开 --eval_all，也输出一个“单 batch”对齐的整体 RMSE 以免落差太大
|
||||
if not args.eval_all:
|
||||
steps = args.s_out
|
||||
xb, yb = X_all[:args.batch], Y_all[:args.batch]
|
||||
y_pred_raw = rollout_raw(xb, steps)
|
||||
y_true_raw = to_raw(yb)
|
||||
y_pred_np = y_pred_raw.cpu().numpy()
|
||||
y_true_np = y_true_raw.cpu().numpy()
|
||||
n1 = np.clip(np.linalg.norm(y_pred_np, axis=-1, keepdims=True), 1e-6, None)
|
||||
n2 = np.clip(np.linalg.norm(y_true_np, axis=-1, keepdims=True), 1e-6, None)
|
||||
y_pred_u = y_pred_np / n1; y_true_u = y_true_np / n2
|
||||
pa = np.arctan2(y_pred_u[...,0], y_pred_u[...,1])
|
||||
ta = np.arctan2(y_true_u[...,0], y_true_u[...,1])
|
||||
diff = wrap_to_pi(pa - ta)
|
||||
rmse = math.sqrt(np.mean(diff**2))
|
||||
print(f"[Test:Mini] angle RMSE = {rmse:.6f} rad ({rmse*180/np.pi:.3f} deg)")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,238 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
启动一个仅 CUDA 的 HTTP 推理服务，输入/输出均为像素坐标。
|
||||
内部将像素轨迹拟合圆并转换为单位 (sin,cos) 以适配训练空间，
|
||||
并把预测结果再还原为像素坐标返回。
|
||||
|
||||
POST /predict
|
||||
body:
|
||||
{
|
||||
"sequences": [ # 批量，每个序列为 [S_in, 2] 像素 [x,y]
|
||||
[[x, y], [x, y], ...],
|
||||
...
|
||||
],
|
||||
"steps": 30, # 预测步数 (>0)
|
||||
"centers": [[cx,cy], ...], # 可选 [B,2]，若已标定则可传入
|
||||
"radii": [r, ...] # 可选 [B]
|
||||
}
|
||||
|
||||
返回：
|
||||
{
|
||||
"pred_xy": [[[x,y], ...], ...], # [B, steps, 2] 像素
|
||||
"device": "cuda",
|
||||
"centers": [[cx,cy], ...], # 用于本次变换的参数
|
||||
"radii": [r, ...]
|
||||
}
|
||||
|
||||
运行示例：
|
||||
uv run serve_pendulum.py --ckpt ckpt/pendulum_tf_base.pt --host 0.0.0.0 --port 8000
|
||||
"""
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
import uvicorn
|
||||
# 兼容直接脚本运行与包内相对导入
|
||||
try:
|
||||
from .geom_utils import batch_pixels_to_unit_sincos, unit_sincos_to_pixels # type: ignore
|
||||
except Exception: # noqa
|
||||
from geom_utils import batch_pixels_to_unit_sincos, unit_sincos_to_pixels
|
||||
|
||||
# ----------------- 与训练一致的模块 -----------------
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model: int, max_len: int = 4096):
|
||||
super().__init__()
|
||||
pe = torch.zeros(max_len, d_model)
|
||||
pos = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
|
||||
div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model))
|
||||
pe[:, 0::2] = torch.sin(pos * div)
|
||||
pe[:, 1::2] = torch.cos(pos * div)
|
||||
self.register_buffer("pe", pe.unsqueeze(0)) # [1, L, D]
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.pe[:, : x.size(1), :]
|
||||
|
||||
class TimeSeriesTransformer(nn.Module):
|
||||
def __init__(self, in_dim=2, d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512, dropout=0.1, out_dim=2):
|
||||
super().__init__()
|
||||
self.in_proj = nn.Linear(in_dim, d_model)
|
||||
self.out_proj = nn.Linear(d_model, out_dim)
|
||||
self.tgt_proj = nn.Linear(out_dim, d_model)
|
||||
self.pos_enc = PositionalEncoding(d_model)
|
||||
self.tf = nn.Transformer(
|
||||
d_model=d_model,
|
||||
nhead=nhead,
|
||||
num_encoder_layers=enc_layers,
|
||||
num_decoder_layers=dec_layers,
|
||||
dim_feedforward=d_ff,
|
||||
dropout=dropout,
|
||||
batch_first=True,
|
||||
)
|
||||
|
||||
def forward(self, src, tgt_in, tgt_mask=None):
|
||||
src = self.pos_enc(self.in_proj(src))
|
||||
tgt = self.pos_enc(self.tgt_proj(tgt_in))
|
||||
out = self.tf(src, tgt, tgt_mask=tgt_mask) # [B, S_out, D]
|
||||
return self.out_proj(out)
|
||||
|
||||
# ----------------- 推理引擎（仅 CUDA） -----------------
|
||||
class InferenceEngine:
|
||||
def __init__(self, ckpt_path: str):
|
||||
if not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA 不可用：该服务仅支持 CUDA。")
|
||||
self.device = torch.device("cuda")
|
||||
|
||||
if not os.path.exists(ckpt_path):
|
||||
raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
|
||||
|
||||
# 显式关闭 weights_only 的安全限制（信任本地 ckpt）
|
||||
payload = torch.load(ckpt_path, map_location=self.device, weights_only=False)
|
||||
|
||||
default_cfg = dict(d_model=256, nhead=4, enc_layers=3, dec_layers=3, d_ff=512)
|
||||
ckpt_cfg = payload.get("cfg", default_cfg)
|
||||
self.net = TimeSeriesTransformer(in_dim=2, out_dim=2, **ckpt_cfg).to(self.device)
|
||||
self.net.load_state_dict(payload["model"], strict=True)
|
||||
self.net.eval()
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
# 标准化参数
|
||||
self.mu = payload["mu"] # numpy
|
||||
self.std = payload["std"]
|
||||
self.mu_t = torch.from_numpy(self.mu).to(self.device)
|
||||
self.std_t = torch.from_numpy(self.std).to(self.device)
|
||||
|
||||
@staticmethod
|
||||
def _unitize_t(t: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
|
||||
n = torch.clamp(torch.linalg.norm(t, dim=-1, keepdim=True), min=eps)
|
||||
return t / n
|
||||
|
||||
def _to_raw(self, z: torch.Tensor) -> torch.Tensor:
|
||||
return z * self.std_t + self.mu_t
|
||||
|
||||
def _to_std(self, z: torch.Tensor) -> torch.Tensor:
|
||||
return (z - self.mu_t) / self.std_t
|
||||
|
||||
def rollout(self, xb_raw: torch.Tensor, steps: int) -> torch.Tensor:
|
||||
"""
|
||||
xb_raw: [B, S_in, 2]，raw 空间（即训练时用的 (sin, cos) 或等价的 [x,y] 单位向量）
|
||||
返回：raw 空间预测 [B, steps, 2]
|
||||
"""
|
||||
xb_raw = self._unitize_t(xb_raw) # 规整到单位圆
|
||||
xb_std = self._to_std(xb_raw) # -> 标准化
|
||||
|
||||
tgt_tokens = xb_std[:, -1:, :2].clone() # 起始 token
|
||||
for _ in range(steps):
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(tgt_tokens.size(1)).to(self.device)
|
||||
out = self.net(xb_std, tgt_tokens, tgt_mask=mask) # 标准化空间
|
||||
next_std = out[:, -1:, :]
|
||||
next_raw = self._to_raw(next_std) # -> raw
|
||||
next_raw = self._unitize_t(next_raw) # 单位圆约束
|
||||
next_std = self._to_std(next_raw) # -> 标准化
|
||||
tgt_tokens = torch.cat([tgt_tokens, next_std], dim=1)
|
||||
|
||||
pred_std = tgt_tokens[:, 1:, :]
|
||||
pred_raw = self._to_raw(pred_std)
|
||||
pred_raw = self._unitize_t(pred_raw) # 返回前再规整
|
||||
return pred_raw
|
||||
|
||||
# ----------------- HTTP 层（仅 /predict） -----------------
|
||||
class PredictRequest(BaseModel):
|
||||
sequences: List[List[List[float]]] = Field(
|
||||
..., description="批量序列：每个为 [S_in,2] 的像素 [x,y] 轨迹（来自图像跟踪）"
|
||||
)
|
||||
steps: int = Field(..., gt=0, description="预测步数 (>0)")
|
||||
# 可选：若客户端已标定可传入，以避免重复拟合
|
||||
centers: Optional[List[List[float]]] = Field(
|
||||
default=None, description="可选 [B,2]，每条序列的圆心 (cx,cy)"
|
||||
)
|
||||
radii: Optional[List[float]] = Field(
|
||||
default=None, description="可选 [B]，每条序列的半径 r"
|
||||
)
|
||||
|
||||
class PredictResponse(BaseModel):
|
||||
pred_xy: List[List[List[float]]] # 未来像素坐标 [B, steps, 2]
|
||||
device: str
|
||||
centers: Optional[List[List[float]]] = None
|
||||
radii: Optional[List[float]] = None
|
||||
|
||||
app = FastAPI(title="Pendulum Transformer Inference (CUDA Only)", version="1.0.0")
|
||||
|
||||
# CORS: 允许任意来源
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=False,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
ENGINE: InferenceEngine | None = None # 运行时注入
|
||||
|
||||
@app.post("/predict", response_model=PredictResponse)
|
||||
def predict(req: PredictRequest):
|
||||
if ENGINE is None:
|
||||
raise HTTPException(500, "Engine not initialized")
|
||||
|
||||
try:
|
||||
pix = np.array(req.sequences, dtype=np.float32) # [B, S_in, 2]
|
||||
except Exception as e:
|
||||
raise HTTPException(400, f"Invalid 'sequences': {e}")
|
||||
|
||||
if pix.ndim != 3 or pix.shape[-1] != 2:
|
||||
raise HTTPException(400, f"'sequences' must be [B, S_in, 2], got {pix.shape}")
|
||||
if req.steps <= 0:
|
||||
raise HTTPException(400, "'steps' must be > 0")
|
||||
if pix.shape[1] < 1:
|
||||
raise HTTPException(400, "Each sequence must have S_in >= 1")
|
||||
|
||||
# 1) 像素 -> 单位 (sin,cos)，并获取每条序列的 (cx,cy,r)
|
||||
if req.centers is not None and req.radii is not None:
|
||||
# 使用来自拍端的标定参数
|
||||
centers = np.array(req.centers, dtype=np.float32) # [B,2]
|
||||
radii = np.array(req.radii, dtype=np.float32) # [B]
|
||||
if centers.shape[0] != pix.shape[0] or centers.shape[1] != 2 or radii.shape[0] != pix.shape[0]:
|
||||
raise HTTPException(400, "centers 应为 [B,2] 且 radii 为 [B]")
|
||||
radii = np.clip(radii, 1e-6, None)
|
||||
# 根据训练时定义：sin=(x-cx)/r, cos=(cy-y)/r
|
||||
sin = (pix[..., 0] - centers[:, 0:1]) / radii[:, None]
|
||||
cos = (centers[:, 1:2] - pix[..., 1]) / radii[:, None]
|
||||
sc = np.stack([sin, cos], axis=-1).astype(np.float32) # [B,S,2]
|
||||
else:
|
||||
sc, centers, radii = batch_pixels_to_unit_sincos(pix)
|
||||
|
||||
# 2) 送入引擎：期望输入为 raw (即单位向量的 sin,cos)
|
||||
x = torch.from_numpy(sc).to(ENGINE.device)
|
||||
pred_sc = ENGINE.rollout(x, steps=req.steps) # [B, steps, 2] 单位 (sin,cos)
|
||||
pred_sc_np = pred_sc.detach().cpu().numpy()
|
||||
|
||||
# 3) 单位 (sin,cos) -> 像素
|
||||
pred_pix = unit_sincos_to_pixels(pred_sc_np, centers, radii) # [B, steps, 2]
|
||||
|
||||
return PredictResponse(
|
||||
pred_xy=pred_pix.tolist(), device="cuda",
|
||||
centers=centers.tolist(), radii=radii.tolist()
|
||||
)
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--ckpt", type=str, required=True, help="Path to checkpoint (.pt)")
|
||||
parser.add_argument("--host", type=str, default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
args = parser.parse_args()
|
||||
|
||||
global ENGINE
|
||||
ENGINE = InferenceEngine(args.ckpt)
|
||||
print(f"[serve] Loaded ckpt from {args.ckpt} on device=cuda")
|
||||
|
||||
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user