lnp_ml/tests/check_gates.py

46 lines
1.6 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import glob
import os
import numpy as np
import torch
BASE = os.environ.get("GATES_BASE", "models/abl")
VARIANTS = os.environ.get("GATES_VARIANTS", "moe,llm,both").split(",")
TASKS = ["size", "delivery", "pdi", "ee", "toxic", "biodist"]
def latest(v: str):
dirs = sorted(glob.glob(f"{BASE}/{v}/*/"), key=os.path.getmtime)
return dirs[-1] if dirs else None
for v in VARIANTS:
run = latest(v)
if run is None:
print(f"\n=== {v}: 无 run 目录 ===")
continue
files = sorted(glob.glob(os.path.join(run, "**", "model.pt"), recursive=True))
print(f"\n=== {v} ({len(files)} folds) ===")
if not files:
print(" 未找到 model.pt本次 run 未保存权重 -> 无法读取门控/log_vars")
continue
gm, gl = [], []
for f in files:
sd = torch.load(f, map_location="cpu", weights_only=False)["model_state_dict"]
for key, acc in (("fusion.g_moe", gm), ("fusion.g_llm", gl)):
if key not in sd:
continue
t = sd[key]
acc.append(abs(t.item()) if t.dim() == 0
else float(t.norm()) / (t.numel() ** 0.5))
for name, acc in (("g_moe", gm), ("g_llm", gl)):
if acc:
print(f" {name}: mean={np.mean(acc):.4f} range=[{min(acc):.4f},{max(acc):.4f}]")
sd = torch.load(files[0], map_location="cpu", weights_only=False)["model_state_dict"]
for t in TASKS:
lv = sd.get(f"head.log_vars.{t}")
if lv is not None:
print(f" {t:9s} log_var={lv.item():+.3f} eff_w={np.exp(-lv.item()):.3f}")