"""主动自我训练（用户数据 → 拆 train/test → 精英保留 + 模型回滚） 策略： 1. 每轮同时采样模型生成 + 随机生成候选 2. 在 train 集上评估 R²，跨轮保留精英候选（elite pool） 3. 只对 R² > 0.1 的正向候选做伪标签 CE 微调（避免强化错误模式） 4. REINFORCE 策略梯度用于模型采样的候选 5. 保存每轮开始前的模型 state，如果本轮最佳 R² 下降则回滚 6. 终止条件：R² > 0.995 或 连续 4 轮 test R² 无改善 """ import warnings import sys import numpy as np import torch import torch.nn.functional as F import config from vocab import Vocab from model import SymbolicRegressionModel from data_generator import RandomExpressionGenerator # 关键：stdout 行缓冲，保证 print() 立即刷新。 # 否则在后台/管道运行时（python 默认块缓冲），进度会攒在一旁， # 直到进程结束才一次性输出 —— 这正是"训练过程不输出"的根因。 if hasattr(sys.stdout, "reconfigure"): sys.stdout.reconfigure(line_buffering=True) # 全局静音候选表达式在数值评估时的 numpy 警告（exp/sqrt/log/power 溢出、 # 定义域外、除零等均属常规现象，结果已用 isfinite/complex 过滤）。 # 全局设置可覆盖 lambdify 生成的已编译模块，避免每一轮都刷屏。 np.seterr(over="ignore", invalid="ignore", divide="ignore") warnings.filterwarnings("ignore", module="lambdifygenerated") COPIED_MODEL_STATE_KEY = "_saved_lstm_state" def train_test_split_arrays(X, y, test_size=0.3, seed=None): """手写随机 train/test 拆分。""" rng = np.random.RandomState(seed) n = X.shape[0] idx = rng.permutation(n) n_te = max(1, int(test_size * n)) te_idx, tr_idx = idx[:n_te], idx[n_te:] return X[tr_idx], X[te_idx], y[tr_idx], y[te_idx] def eval_expr_on_data(expr_str: str, X: np.ndarray, y_true: np.ndarray) -> tuple: """计算给定表达式在数据上的预测值和 MSE。无效返回 NaN。""" import sympy as sp import warnings D = X.shape[1] try: expr = sp.sympify(expr_str) symbols = [sp.Symbol(f"x{i}") for i in range(D)] f = sp.lambdify(symbols, expr, modules="numpy") # 候选表达式常落在定义域外或数值溢出（exp/sqrt/log/power）， # 静音这些 numpy RuntimeWarning，避免刷屏。结果 NaN/inf 会被下方过滤。 with warnings.catch_warnings(): warnings.simplefilter("ignore") with np.errstate(over="ignore", invalid="ignore", divide="ignore"): y_pred = f(*[X[:, i] for i in range(D)]) if np.iscomplexobj(y_pred) or not np.isfinite(y_pred).all(): return np.full_like(y_true, np.nan), np.nan mse = float(np.nanmean((y_pred - y_true) ** 2)) return y_pred, mse except Exception: return np.full_like(y_true, np.nan), np.nan def _generate_random_candidates(num_candidates: int, D: int, seed: int = None) -> list[str]: """生成随机候选，混合深度。""" gen = RandomExpressionGenerator(seed=seed) rng = np.random.RandomState(seed) candidates = [] depths = [2, 3, 4] while len(candidates) < num_candidates * 3: try: depth = rng.choice(depths) data = gen.generate(var_count=D, max_depth=int(depth)) candidates.append(data["expr_str"]) except Exception: continue return candidates[:num_candidates] def _evaluate_candidates(expr_strings: list[str], X: np.ndarray, y: np.ndarray, y_var: float) -> tuple: """批量评估候选，返回 rewards (R²) 和 mses。无效 reward=-1。""" n = len(expr_strings) rewards = np.full(n, -1.0) mses = np.full(n, np.nan) for i, expr_str in enumerate(expr_strings): _, mse = eval_expr_on_data(expr_str, X, y) mses[i] = mse if np.isfinite(mse): rewards[i] = 1.0 - mse / max(y_var, 1e-9) return rewards, mses @torch.no_grad() def _encode_input(X: np.ndarray, y: np.ndarray, device: str): """编码用户数据为模型输入 [1, N, D_MAX], [1, N].""" D = X.shape[1] N = X.shape[0] X_t = torch.zeros(1, N, config.D_MAX, dtype=torch.float32, device=device) X_t[0, :, :D] = torch.from_numpy(X).to(device) y_t = torch.from_numpy(y).float().unsqueeze(0).to(device) n_samples = torch.tensor([N], device=device) var_counts = torch.tensor([D], device=device) return X_t, y_t, n_samples, var_counts @torch.enable_grad() def _sample_model(model, X_t, y_t, n_samples, var_counts, device, max_len=None, num_candidates=20, greedy=False, temperature=1.5): """从模型采样候选。greedy → argmax；否则多项式采样 + 噪声扰动。""" if max_len is None: max_len = config.T_MAX sample_size = num_candidates vocab = model.vocab z = model.encoder(X_t, y_t, n_samples, var_counts) if not greedy: noise_scale = 0.15 z = z + torch.randn_like(z) * noise_scale temp = temperature else: temp = 0.1 hidden = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous() cell = z.unsqueeze(0).expand(-1, sample_size, -1).contiguous() if not greedy: hidden = hidden + torch.randn_like(hidden) * noise_scale cell = cell + torch.randn_like(cell) * noise_scale token = torch.full((sample_size, 1), fill_value=vocab.sos_id, device=device) sequences = [token] log_probs = [] done = torch.zeros(sample_size, dtype=torch.bool, device=device) for _ in range(max_len): embedded = model.lstm.embedding(token) output, (hidden, cell) = model.lstm.lstm(embedded, (hidden, cell)) logits = model.lstm.proj(output).squeeze(1) logits = logits / temp if greedy: token = logits.argmax(dim=-1, keepdim=True) else: log_probs_t = F.log_softmax(logits, dim=-1) probs_t = F.softmax(logits, dim=-1) dist = torch.distributions.Categorical(probs=probs_t) token = dist.sample().unsqueeze(-1) lp = torch.gather(log_probs_t, 1, token) lp = lp.masked_fill(done.unsqueeze(-1), 0.0) log_probs.append(lp) sequences.append(token) done = done | (token.squeeze(-1) == vocab.eos_id) if done.all(): break tokens = torch.cat(sequences, dim=1) if not greedy: log_probs = torch.cat(log_probs, dim=1) tokens = tokens[:, 1:] expr_strings = [] for i in range(sample_size): toks = tokens[i] if (toks == vocab.eos_id).any(): eos = (toks == vocab.eos_id).nonzero(as_tuple=True)[0][0].item() toks = toks[:eos] expr_strings.append(vocab.decode(toks.tolist())) return tokens, log_probs if not greedy else None, expr_strings def _save_model_snapshot(model): """保存 LSTM 参数快照到 model 的临时属性。""" sd = {k: v.data.clone() for k, v in model.lstm.state_dict().items()} setattr(model, COPIED_MODEL_STATE_KEY, sd) def _restore_model_snapshot(model): """从快照恢复 LSTM 参数。""" sd = getattr(model, COPIED_MODEL_STATE_KEY, None) if sd is not None: for k, v in sd.items(): model.lstm.state_dict()[k].data.copy_(v) def _encode_expr(expr_str: str, vocab: Vocab, device: str): """把表达式字符串编码为 token 序列。""" token_ids = vocab.encode(expr_str) token_ids = [vocab.sos_id] + token_ids + [vocab.eos_id] token_ids = token_ids[:config.T_MAX] return torch.tensor(token_ids, dtype=torch.long, device=device).unsqueeze(0) def self_train_on_user_data(model, X, y, device, split=0.7, max_rounds=15, candidates=96, seed=None, verbose=True): """针对用户数据 (X, y) 执行主动自我训练。 返回 dict {best_expr, best_r2, converged, rounds, results, round_results} """ vocab = model.vocab if X.ndim == 1: X = X.reshape(-1, 1) D = X.shape[1] Xtr, Xte, ytr, yte = train_test_split_arrays( X, y, test_size=1 - split, seed=seed ) ytr_var = float(np.var(ytr)) yte_var = float(np.var(yte)) if verbose: print(f" [数据] 训练集: {Xtr.shape[0]} 样本, 测试集: {Xte.shape[0]} 样本, 维度 D={D}") print(f" [数据] y 方差: train={ytr_var:.4f}, test={yte_var:.4f}") Xtr_t, ytr_t, ntr, vctr = _encode_input(Xtr, ytr, device) optimizer = torch.optim.Adam(model.lstm.parameters(), lr=3e-4) # ---- 精英池：跨轮保留历史最佳候选 ---- elite_pool = [] PATIENCE = 4 results = [] best_r2 = -float("inf") best_expr = None no_improve = 0 converged = False for rnd in range(1, max_rounds + 1): model.train() # --- 保存本轮开始前的模型 --- _save_model_snapshot(model) # --- 温度退火 --- temp = max(1.2, 3.0 - 0.18 * (rnd - 1)) # --- 生成候选 --- n_model = int(candidates * 0.45) n_rand = candidates - n_model # 模型采样 tokens, log_probs, model_exprs = _sample_model( model, Xtr_t, ytr_t, ntr, vctr, device, num_candidates=n_model, greedy=False, temperature=temp ) # 随机生成（混合深度，提供多样性） rand_candidates = _generate_random_candidates( n_rand, D, seed=(seed or 0) + rnd * 131 + 7 ) # 加入精英池中的候选（去重后按 R² 排序取前10） elite_exprs = [e for e, _ in elite_pool[:10]] expr_strings = list(model_exprs) + list(rand_candidates) + elite_exprs # --- 评估候选 --- rewards, mses = _evaluate_candidates(expr_strings, Xtr, ytr, ytr_var) valid_mask = np.isfinite(mses) valid_idxs = np.where(valid_mask)[0] if valid_idxs.size == 0: if verbose: print(f" [自我训练] round {rnd}/{max_rounds} | T={temp:.2f} | 所有候选无效，跳过此轮") results.append({"round": rnd, "train_r2": -1.0, "test_r2": -1.0, "test_expr": "", "n_valid": 0}) continue best_local_idx = valid_idxs[int(np.argmax(rewards[valid_idxs]))] best_local_expr = expr_strings[best_local_idx] best_local_r2 = float(rewards[best_local_idx]) # --- 更新精英池 --- for i in valid_idxs: r2 = float(rewards[i]) if r2 > 0.05: elite_pool.append((expr_strings[i], r2)) elite_pool = sorted(set(elite_pool), key=lambda x: -x[1])[:20] # --- 训练步：伪标签 + REINFORCE --- # 只对 R² > 0.1 的正向候选做 CE 伪标签 good_mask = rewards > 0.1 good_idxs = np.where(good_mask)[0] ce_losses = [] for idx in good_idxs: expr_str = expr_strings[idx] if expr_str == "": continue try: t_tensor = _encode_expr(expr_str, vocab, device) z = model.encoder(Xtr_t, ytr_t, ntr, vctr) hidden = z.unsqueeze(0) cell = z.unsqueeze(0) logits, _, _ = model.lstm(t_tensor[:, :-1], hidden, cell) targets = t_tensor[:, 1:] ce = F.cross_entropy(logits.reshape(-1, vocab.vocab_size), targets.reshape(-1)) ce_losses.append(ce * 0.3) except Exception: continue if ce_losses and verbose: pos_count = len(ce_losses) best_r2_among_pos = max(rewards[good_idxs]) print(f" └─ 伪标签: {pos_count} 个正向候选 (R²>{0.1:.1f}, 最高={best_r2_among_pos:.3f})") # REINFORCE：仅对模型采样候选（前 n_model 个）中有效且导通的 reinforce_loss = None model_valid_idx = valid_idxs[valid_idxs < n_model] non_collapsed = [] for idx in model_valid_idx: if idx < n_model and log_probs is not None: lp_row = log_probs[idx] if lp_row.numel() > 0 and lp_row.sum() < -0.01: non_collapsed.append(idx) if len(non_collapsed) >= 3: m_rewards = rewards[non_collapsed] m_rewards = np.where(np.isfinite(m_rewards), m_rewards, -1.0) m_advantage_t = torch.tensor(m_rewards - m_rewards.mean(), dtype=torch.float32, device=device) m_logp = log_probs[non_collapsed] m_logp_sum = m_logp.sum(dim=-1) reinforce_loss = -(m_advantage_t * m_logp_sum).mean() if ce_losses or reinforce_loss is not None: all_losses = ce_losses[:] if reinforce_loss is not None: all_losses.append(reinforce_loss * 0.5) loss = torch.stack(all_losses).mean() optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.lstm.parameters(), 1.0) optimizer.step() # --- test 评估 --- te_expr = best_local_expr _, te_mse = eval_expr_on_data(te_expr, Xte, yte) if np.isfinite(te_mse): te_r2 = 1.0 - te_mse / max(yte_var, 1e-9) else: te_r2 = -float("inf") # --- 模型回滚：如果本轮最佳 R² 比上一轮差，恢复模型 --- prev_r2 = results[-1]["best_local_r2"] if results else -float("inf") if best_local_r2 < prev_r2 - 0.1 and len(results) >= 1: _restore_model_snapshot(model) if verbose: print(f" └─ 模型回滚（best R² {best_local_r2:.3f} < 上一轮 {prev_r2:.3f}）") # 重新用恢复后的模型采样一次作为补充 tokens2, lp2, model_exprs2 = _sample_model( model, Xtr_t, ytr_t, ntr, vctr, device, num_candidates=n_model, greedy=False, temperature=temp ) expr_strings2 = list(model_exprs2) + elite_exprs[:5] rewards2, mses2 = _evaluate_candidates(expr_strings2, Xtr, ytr, ytr_var) vm2 = np.isfinite(mses2) vi2 = np.where(vm2)[0] if vi2.size > 0: best_i2 = vi2[int(np.argmax(rewards2[vi2]))] best_local_expr = expr_strings2[best_i2] best_local_r2 = float(rewards2[best_i2]) _, te_mse2 = eval_expr_on_data(best_local_expr, Xte, yte) if np.isfinite(te_mse2): te_r2 = 1.0 - te_mse2 / max(yte_var, 1e-9) else: te_r2 = -float("inf") if te_r2 > best_r2: best_r2 = te_r2 best_expr = te_expr no_improve = 0 else: no_improve += 1 nv = int(valid_mask.sum()) results.append({ "round": rnd, "train_r2": float(np.nanmean(rewards)), "test_r2": float(te_r2), "test_expr": te_expr, "best_local_expr": best_local_expr, "best_local_r2": best_local_r2, "n_valid": nv, "n_good_ce": len(ce_losses), }) if verbose: avg_r2 = float(np.nanmean(rewards)) print(f" [自我训练] round {rnd}/{max_rounds} | T={temp:.2f} | " f"valid={nv}/{candidates} | " f"avg R²={avg_r2:.4f} | " f"best R²={best_local_r2:.4f} | " f"test R²={te_r2:.4f} | {te_expr}") # --- 终止判定 --- if te_r2 > 0.995: if verbose: print(f" → R² 达标 ({te_r2:.4f})，停止自我训练 ✓") converged = True break if no_improve >= PATIENCE: if verbose: print(f" → test R² 连续 {PATIENCE} 轮无改善，停止自我训练（最佳: {best_r2:.4f}）") converged = True break return { "best_expr": best_expr, "best_r2": best_r2, "converged": converged, "rounds": len(results), "round_results": results, }