rpotter6298
113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
"""inspect_embeddings — per-dim statistics on a run's saved feature embeddings.
|
|
|
|
Loads features.h5 from a run folder (must have been produced with
|
|
``save_features: true`` in the config), computes per-dimension variance,
|
|
sparsity (fraction of |z| < threshold), and useful summary stats.
|
|
|
|
Helpful for diagnosing whether a fusion stage is collapsing dimensions to
|
|
near-zero — which would silently null out information when the bridge uses
|
|
a Hadamard product.
|
|
|
|
Usage:
|
|
python -m v4.scripts.analysis.inspect_embeddings <run_folder> [--phase NAME]
|
|
[--threshold 0.05]
|
|
|
|
Examples:
|
|
python -m v4.scripts.analysis.inspect_embeddings \
|
|
v4/results/experiments/tri_v1/baseline_tri/rep00 --phase nt
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import h5py
|
|
import numpy as np
|
|
|
|
|
|
def load_phase(features_path: Path, phase: str | None):
|
|
with h5py.File(features_path, "r") as f:
|
|
phases = list(f.keys())
|
|
if phase is None:
|
|
phase = phases[-1]
|
|
if phase not in phases:
|
|
raise SystemExit(
|
|
f"Phase {phase!r} not in {features_path} (available: {phases})"
|
|
)
|
|
g = f[phase]
|
|
emb_keys = [k for k in g.keys() if k.endswith("_embedding")]
|
|
if not emb_keys:
|
|
raise SystemExit(
|
|
f"No '*_embedding' dataset in phase {phase!r} of {features_path}"
|
|
)
|
|
# Prefer the embedding named after the phase if present, else first match.
|
|
emb_key = f"{phase}_embedding" if f"{phase}_embedding" in emb_keys else emb_keys[0]
|
|
z_raw = g[emb_key][:]
|
|
# Stored shape is (n_folds, n_epochs, n_samples, n_dim). Use last epoch.
|
|
if z_raw.ndim == 4:
|
|
z = z_raw[:, -1, :, :]
|
|
else:
|
|
z = z_raw
|
|
split = g["split"][:].astype(str)
|
|
y_true = g["y_true"][:]
|
|
return phase, z, split, y_true, phases
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("run_dir", type=Path, help="Folder containing features.h5 (under repNN)")
|
|
ap.add_argument("--phase", default=None,
|
|
help="Bridge phase name (e.g. nt, hb, cd_fuse). Default: last in file.")
|
|
ap.add_argument("--threshold", type=float, default=0.05,
|
|
help="|z| threshold for the sparsity count (default: 0.05)")
|
|
args = ap.parse_args()
|
|
|
|
feat = next(iter(args.run_dir.rglob("features.h5")), None)
|
|
if feat is None:
|
|
raise SystemExit(f"No features.h5 found under {args.run_dir}")
|
|
|
|
phase, z, split, _, all_phases = load_phase(feat, args.phase)
|
|
|
|
val_mask = (split == "val")
|
|
z_val = z[val_mask]
|
|
z_train = z[(split == "train")]
|
|
n_dim = z_val.shape[-1] if z_val.size else 0
|
|
if n_dim == 0:
|
|
raise SystemExit(f"No val embeddings found for phase {phase!r}")
|
|
|
|
print(f"Run: {args.run_dir}")
|
|
print(f"File: {feat}")
|
|
print(f"Phases: {all_phases}")
|
|
print(f"Phase: {phase} (z shape: {z.shape})")
|
|
print(f"Split sizes: train={len(z_train)} val={len(z_val)}")
|
|
print()
|
|
|
|
abs_z = np.abs(z_val)
|
|
per_dim_var = z_val.var(axis=0)
|
|
per_dim_abs = abs_z.mean(axis=0)
|
|
per_dim_max = abs_z.max(axis=0)
|
|
sparsity = (abs_z < args.threshold).mean(axis=0)
|
|
|
|
print(f"Aggregate stats over val embeddings (|z| < {args.threshold} = 'near-zero'):")
|
|
print(f" global mean(|z|): {abs_z.mean():.4f}")
|
|
print(f" global var(z): {z_val.var():.4f}")
|
|
print(f" fraction near-zero (global): {(abs_z < args.threshold).mean():.4f}")
|
|
print()
|
|
print(f"Per-dimension summary ({n_dim} dims):")
|
|
print(f" variance: min={per_dim_var.min():.4f} med={np.median(per_dim_var):.4f} max={per_dim_var.max():.4f}")
|
|
print(f" mean |z|: min={per_dim_abs.min():.4f} med={np.median(per_dim_abs):.4f} max={per_dim_abs.max():.4f}")
|
|
print(f" near-zero rate: min={sparsity.min():.4f} med={np.median(sparsity):.4f} max={sparsity.max():.4f}")
|
|
|
|
# Dead dimensions: high near-zero rate
|
|
dead = np.where(sparsity > 0.9)[0]
|
|
print(f" 'dead' dims (>90% near-zero): {len(dead)}/{n_dim} "
|
|
f"{('idx: ' + str(dead.tolist())) if 0 < len(dead) <= 20 else ''}")
|
|
weak = np.where(per_dim_var < 1e-4)[0]
|
|
print(f" 'weak' dims (var < 1e-4): {len(weak)}/{n_dim} "
|
|
f"{('idx: ' + str(weak.tolist())) if 0 < len(weak) <= 20 else ''}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|