817 lines
34 KiB
Python
817 lines
34 KiB
Python
#!/usr/bin/env python3
|
||
import os
|
||
import re
|
||
import csv
|
||
import glob
|
||
import argparse
|
||
import torch
|
||
import torch.nn.functional as F
|
||
from collections import defaultdict
|
||
|
||
try:
|
||
import matplotlib
|
||
matplotlib.use("Agg") # headless – no display needed
|
||
import matplotlib.pyplot as plt
|
||
import matplotlib.ticker as ticker
|
||
_MATPLOTLIB_OK = True
|
||
except ImportError:
|
||
_MATPLOTLIB_OK = False
|
||
|
||
try:
|
||
from safetensors.torch import load_file
|
||
except ImportError:
|
||
print("Error: safetensors must be installed. Please run `pip install safetensors`.")
|
||
exit(1)
|
||
|
||
|
||
# ── filename parsing ──────────────────────────────────────────────────────────
|
||
|
||
FILENAME_RE = re.compile(r"^(\d+)_S(\d+)(?:_(uncond))?\.safetensors$")
|
||
|
||
def parse_filename(filename: str):
|
||
"""
|
||
Returns (seed, step, is_uncond) or None if the filename doesn't match.
|
||
Expected formats:
|
||
<seed>_S<step>.safetensors → cond
|
||
<seed>_S<step>_uncond.safetensors → uncond
|
||
"""
|
||
m = FILENAME_RE.match(filename)
|
||
if not m:
|
||
return None
|
||
seed = m.group(1)
|
||
step = int(m.group(2))
|
||
is_uncond = m.group(3) is not None
|
||
return seed, step, is_uncond
|
||
|
||
|
||
# ── per-tensor metrics ────────────────────────────────────────────────────────
|
||
|
||
def compute_metrics(bl: torch.Tensor, qt: torch.Tensor) -> dict:
|
||
"""
|
||
All metrics between a paired (baseline, quantized) latent tensor.
|
||
Both tensors are float32, same shape.
|
||
"""
|
||
diff = bl - qt # signed error
|
||
# ── basic error ──
|
||
mse = (diff ** 2).mean().item()
|
||
mae = diff.abs().mean().item()
|
||
max_err = diff.abs().max().item()
|
||
|
||
# ── relative error (avoid div-by-zero) ──
|
||
bl_norm = (bl ** 2).mean().sqrt().item()
|
||
rel_rmse = (mse ** 0.5) / (bl_norm + 1e-8)
|
||
|
||
# ── SNR in dB: signal power / noise power ──
|
||
signal_pow = (bl ** 2).mean().item()
|
||
noise_pow = mse
|
||
snr_db = 10.0 * torch.log10(
|
||
torch.tensor(signal_pow / (noise_pow + 1e-12))
|
||
).item()
|
||
|
||
# ── cosine similarity (flatten spatial dims, keep channel dim) ──
|
||
# shape: (C, H*W) → per-channel cosine, then average
|
||
c = bl.shape[0] if bl.dim() >= 1 else 1
|
||
bl_flat = bl.reshape(c, -1)
|
||
qt_flat = qt.reshape(c, -1)
|
||
cos_sim = F.cosine_similarity(bl_flat, qt_flat, dim=1).mean().item()
|
||
|
||
# ── per-channel MSE stats ──
|
||
per_ch_mse = (diff.reshape(c, -1) ** 2).mean(dim=1) # (C,)
|
||
ch_mse_mean = per_ch_mse.mean().item()
|
||
ch_mse_std = per_ch_mse.std().item() if per_ch_mse.numel() > 1 else 0.0 #per_ch_mse.std().item()
|
||
ch_mse_max = per_ch_mse.max().item()
|
||
worst_channel = per_ch_mse.argmax().item()
|
||
|
||
# ── structural: variance preservation ──
|
||
# How well does the quantized model preserve the spatial variance?
|
||
bl_var = bl.var().item()
|
||
qt_var = qt.var().item()
|
||
var_ratio = qt_var / (bl_var + 1e-12) # 1.0 = perfect
|
||
|
||
# ── outlier fraction: pixels where |error| > 2σ of baseline ──
|
||
threshold = 2.0 * bl.std().item()
|
||
outlier_frac = (diff.abs() > threshold).float().mean().item()
|
||
|
||
return {
|
||
"mse": mse,
|
||
"mae": mae,
|
||
"max_err": max_err,
|
||
"rel_rmse": rel_rmse,
|
||
"snr_db": snr_db,
|
||
"cos_sim": cos_sim,
|
||
"ch_mse_mean": ch_mse_mean,
|
||
"ch_mse_std": ch_mse_std,
|
||
"ch_mse_max": ch_mse_max,
|
||
"worst_channel": worst_channel,
|
||
"var_ratio": var_ratio,
|
||
"outlier_frac": outlier_frac,
|
||
}
|
||
|
||
|
||
# ── marginal (per-step) delta metrics ────────────────────────────────────────
|
||
|
||
def compute_marginal(cur_metrics: dict, prev_metrics: dict) -> dict:
|
||
"""
|
||
Scalar delta between consecutive steps for the key accumulation metrics.
|
||
Positive = error grew this step; negative = error shrank (rare but possible).
|
||
"""
|
||
keys = ("mse", "mae", "snr_db", "cos_sim", "outlier_frac")
|
||
return {f"delta_{k}": cur_metrics[k] - prev_metrics[k] for k in keys}
|
||
|
||
|
||
# ── helpers ───────────────────────────────────────────────────────────────────
|
||
|
||
def _nanstats(vals: list[float]) -> tuple[float, float, int]:
|
||
"""Return (mean, std, count) ignoring NaNs."""
|
||
clean = [v for v in vals if v == v] # NaN != NaN
|
||
n = len(clean)
|
||
if n == 0:
|
||
return float("nan"), float("nan"), 0
|
||
mean = sum(clean) / n
|
||
var = sum((v - mean) ** 2 for v in clean) / n if n > 1 else 0.0
|
||
return mean, var ** 0.5, n
|
||
|
||
|
||
def _winner_indices(col_means: list[float], higher_is_better: bool) -> set[int]:
|
||
"""Return indices of methods that are strictly best for this metric."""
|
||
clean = [(i, v) for i, v in enumerate(col_means) if v == v]
|
||
if not clean:
|
||
return set()
|
||
best = max(clean, key=lambda x: x[1])[1] if higher_is_better \
|
||
else min(clean, key=lambda x: x[1])[1]
|
||
return {i for i, v in clean if v == best}
|
||
|
||
|
||
def _stratified_stats(rows: list[dict], col: str) -> tuple[float, float, int]:
|
||
"""
|
||
Compute mean and SE where the SE is the average of per-(step, cond_type)
|
||
standard errors (std / sqrt(n_seeds)) rather than raw stds.
|
||
|
||
Stratifying by (step, cond_type) removes the huge timestep-to-timestep
|
||
scale variance, leaving only seed-to-seed noise. Dividing by sqrt(n_seeds)
|
||
then converts that noise into a standard error — a much tighter and more
|
||
meaningful uncertainty on the *mean* estimate.
|
||
"""
|
||
by_step: dict[tuple, list[float]] = defaultdict(list)
|
||
for r in rows:
|
||
v = r[col]
|
||
if v == v: # skip NaN
|
||
by_step[(r["step"], r["cond_type"])].append(v)
|
||
|
||
if not by_step:
|
||
return float("nan"), float("nan"), 0
|
||
|
||
step_means, step_ses, total_n = [], [], 0
|
||
for vals in by_step.values():
|
||
n = len(vals)
|
||
if n == 0:
|
||
continue
|
||
m = sum(vals) / n
|
||
std = (sum((v - m) ** 2 for v in vals) / n) ** 0.5 if n > 1 else 0.0
|
||
se = std / (n ** 0.5) # SE per unique seed at this stratum
|
||
step_means.append(m)
|
||
step_ses.append(se)
|
||
total_n += n
|
||
|
||
if not step_means:
|
||
return float("nan"), float("nan"), 0
|
||
|
||
return sum(step_means) / len(step_means), sum(step_ses) / len(step_ses), total_n
|
||
|
||
|
||
def print_comparison_table(
|
||
all_quant_data: dict[str, list[dict]],
|
||
baseline_name: str,
|
||
stratify_std: bool = False,
|
||
) -> None:
|
||
"""
|
||
Print a single condensed table: rows = metrics, columns = quant methods.
|
||
Shows mean ± std. Best value per row is marked with ★.
|
||
|
||
stratify_std=True → std is the avg of per-timestep seed stds.
|
||
Much tighter and more meaningful for scale-variant metrics.
|
||
stratify_std=False → std is pooled over all seeds × steps × cond_types.
|
||
"""
|
||
# ── ANSI colour palette ───────────────────────────────────────────────────
|
||
RESET = "\033[0m"
|
||
BOLD = "\033[1m"
|
||
DIM = "\033[2m"
|
||
|
||
# Row alternating backgrounds (dark-terminal friendly)
|
||
ROW_A = "\033[48;5;234m" # very dark grey
|
||
ROW_B = "\033[48;5;236m" # slightly lighter dark grey
|
||
|
||
# Accent colours
|
||
TITLE_FG = "\033[38;5;117m" # soft cyan
|
||
HEADER_FG = "\033[38;5;153m" # light blue-white
|
||
LABEL_FG = "\033[38;5;252m" # near-white
|
||
VALUE_FG = "\033[38;5;222m" # warm yellow
|
||
STD_FG = "\033[38;5;244m" # mid-grey
|
||
STAR_FG = "\033[38;5;214m" # amber/gold
|
||
BORDER_FG = "\033[38;5;240m" # dim grey
|
||
NOTE_FG = "\033[38;5;102m" # muted green-grey
|
||
|
||
TABLE_METRICS = [
|
||
("mse", "MSE ↓", False, ".5f"),
|
||
("mae", "MAE ↓", False, ".5f"),
|
||
("max_err", "Max err ↓", False, ".5f"),
|
||
("rel_rmse", "Rel-RMSE ↓", False, ".5f"),
|
||
("snr_db", "SNR dB ↑", True, ".2f"),
|
||
("cos_sim", "Cos-sim ↑", True, ".6f"),
|
||
("var_ratio", "Var ratio →1", None, ".5f"), # closest to 1 wins
|
||
("outlier_frac", "Outlier% ↓", False, ".5f"),
|
||
("ch_mse_max", "Ch-MSE max ↓", False, ".5f"),
|
||
("ch_mse_std", "Ch-MSE std ↓", False, ".5f"),
|
||
("delta_mse", "ΔMSE/step ↓", False, ".6f"),
|
||
("delta_cos_sim","ΔCos/step ↑", True, ".7f"),
|
||
]
|
||
|
||
methods = list(all_quant_data.keys())
|
||
if not methods:
|
||
return
|
||
|
||
stats: dict[str, dict[str, tuple]] = {}
|
||
for method, rows in all_quant_data.items():
|
||
stats[method] = {}
|
||
for col, *_ in TABLE_METRICS:
|
||
if stratify_std:
|
||
stats[method][col] = _stratified_stats(rows, col)
|
||
else:
|
||
stats[method][col] = _nanstats([r[col] for r in rows])
|
||
|
||
# ── column widths ─────────────────────────────────────────────────────────
|
||
# A cell looks like: "0.07981 ±0.01105 ★" (value ± std [star])
|
||
# We compute the rendered (no-ANSI) width of the widest possible cell.
|
||
LABEL_W = max(len(lbl) for _, lbl, *_ in TABLE_METRICS)
|
||
# Widest value string per method (scan actual data)
|
||
def _cell_text(method, col, fmt):
|
||
mean, std, n = stats[method][col]
|
||
if n == 0:
|
||
return "N/A"
|
||
return f"{mean:{fmt}} ±{std:{fmt}} ★"
|
||
|
||
METHOD_W = max(len(m) for m in methods)
|
||
VALUE_COL_W = max(
|
||
max(len(_cell_text(m, col, fmt)) for m in methods)
|
||
for col, _, _, fmt in TABLE_METRICS
|
||
)
|
||
CELL_W = max(METHOD_W + 2, VALUE_COL_W, 22) # at least wide enough for header
|
||
|
||
# ── separator geometry ────────────────────────────────────────────────────
|
||
# " " + label + " │" + n×(" " + cell + " │")
|
||
total_w = 2 + LABEL_W + 2 + 1 + len(methods) * (1 + CELL_W + 2 + 1)
|
||
THICK = "═"
|
||
THIN = "─"
|
||
|
||
def _border(left, mid, right, horiz):
|
||
"""Build a box-drawing separator line."""
|
||
seg = horiz * (CELL_W + 2)
|
||
inner = mid.join(seg for _ in methods)
|
||
return f"{left}{horiz * (LABEL_W + 4)}{mid}{inner}{right}"
|
||
|
||
top_border = _border("╔", "╦", "╗", "═")
|
||
hdr_border = _border("╠", "╬", "╣", "═")
|
||
row_div = _border("╟", "╫", "╢", "─")
|
||
bot_border = _border("╚", "╩", "╝", "═")
|
||
|
||
std_label = ("± = avg of per-timestep SE (std/√n_seeds) [--stratify-std]"
|
||
if stratify_std else
|
||
"± = std pooled over all seeds × steps × cond_types")
|
||
|
||
# ── title bar ─────────────────────────────────────────────────────────────
|
||
title_text = f" CONDENSED COMPARISON TABLE (vs baseline: {baseline_name}) "
|
||
title_pad = max(len(top_border), len(title_text) + 4)
|
||
|
||
print()
|
||
print(f"{BOLD}{TITLE_FG}{THICK * title_pad}{RESET}")
|
||
print(f"{BOLD}{TITLE_FG}{'':^{title_pad}}{RESET}")
|
||
centered_title = title_text.center(title_pad)
|
||
print(f"{BOLD}{TITLE_FG}{centered_title}{RESET}")
|
||
print(f"{BOLD}{TITLE_FG}{'':^{title_pad}}{RESET}")
|
||
print(f"{BOLD}{TITLE_FG}{THICK * title_pad}{RESET}")
|
||
|
||
# ── table top border ──────────────────────────────────────────────────────
|
||
print(f"{BORDER_FG}{top_border}{RESET}")
|
||
|
||
# ── column header row ─────────────────────────────────────────────────────
|
||
metric_hdr = f" {'Metric':<{LABEL_W}} "
|
||
col_hdrs = " ".join(f"{BOLD}{HEADER_FG}{m:^{CELL_W}}{RESET}" for m in methods)
|
||
# rebuild with borders
|
||
col_hdrs_b = "".join(
|
||
f"{BORDER_FG}│{RESET} {BOLD}{HEADER_FG}{m:^{CELL_W}}{RESET} "
|
||
for m in methods
|
||
) + f"{BORDER_FG}│{RESET}"
|
||
print(f"{BORDER_FG}║{RESET}{BOLD}{HEADER_FG}{metric_hdr}{RESET}{col_hdrs_b}")
|
||
print(f"{BORDER_FG}{hdr_border}{RESET}")
|
||
|
||
# ── data rows ─────────────────────────────────────────────────────────────
|
||
for row_idx, (col, label, higher_is_better, fmt) in enumerate(TABLE_METRICS):
|
||
means = [stats[m][col][0] for m in methods]
|
||
|
||
if higher_is_better is None:
|
||
clean = [(i, v) for i, v in enumerate(means) if v == v]
|
||
winners = {min(clean, key=lambda x: abs(x[1] - 1.0))[0]} if clean else set()
|
||
else:
|
||
winners = _winner_indices(means, higher_is_better)
|
||
|
||
row_bg = ROW_A if row_idx % 2 == 0 else ROW_B
|
||
|
||
# label cell
|
||
label_cell = f"{row_bg}{LABEL_FG} {label:<{LABEL_W}} {RESET}"
|
||
|
||
# value cells
|
||
value_cells = ""
|
||
for i, method in enumerate(methods):
|
||
mean, std, n = stats[method][col]
|
||
if n == 0:
|
||
inner = f"{'N/A':^{CELL_W}}"
|
||
value_cells += f"{BORDER_FG}│{RESET}{row_bg} {DIM}{inner}{RESET} "
|
||
else:
|
||
val_str = f"{mean:{fmt}}"
|
||
std_str = f"±{std:{fmt}}"
|
||
is_win = i in winners
|
||
star_str = f" {STAR_FG}★{RESET}{row_bg}" if is_win else " "
|
||
|
||
# Build the content string (raw, no ANSI) to measure padding
|
||
raw_content = f"{val_str} {std_str}{' ★' if is_win else ' '}"
|
||
pad_total = CELL_W - len(raw_content)
|
||
lpad = pad_total // 2
|
||
rpad = pad_total - lpad
|
||
|
||
colored = (
|
||
f"{' ' * lpad}"
|
||
f"{BOLD if is_win else ''}{VALUE_FG}{val_str}{RESET}{row_bg} "
|
||
f"{STD_FG}{std_str}{RESET}{row_bg}"
|
||
f"{star_str}"
|
||
f"{' ' * rpad}"
|
||
)
|
||
value_cells += f"{BORDER_FG}│{RESET}{row_bg} {colored}{RESET} "
|
||
|
||
value_cells += f"{BORDER_FG}│{RESET}"
|
||
print(f"{BORDER_FG}║{RESET}{label_cell}{value_cells}")
|
||
|
||
# thin divider between rows (skip after last)
|
||
if row_idx < len(TABLE_METRICS) - 1:
|
||
print(f"{BORDER_FG}{row_div}{RESET}")
|
||
|
||
# ── bottom border & legend ────────────────────────────────────────────────
|
||
print(f"{BORDER_FG}{bot_border}{RESET}")
|
||
print(f" {STAR_FG}★{RESET} {NOTE_FG}= best value for that metric"
|
||
f" │ {std_label}{RESET}")
|
||
print()
|
||
|
||
|
||
# ── github markdown table ─────────────────────────────────────────────────────
|
||
|
||
def save_markdown_table(
|
||
all_quant_data: dict[str, list[dict]],
|
||
baseline_name: str,
|
||
out_path: str,
|
||
stratify_std: bool = False,
|
||
) -> None:
|
||
"""
|
||
Write a GitHub-flavoured Markdown comparison table to *out_path*.
|
||
Mirrors the logic of print_comparison_table but outputs pipe-table syntax.
|
||
Best value per metric is marked with ★.
|
||
"""
|
||
TABLE_METRICS = [
|
||
("mse", "MSE ↓", False, ".5f"),
|
||
("mae", "MAE ↓", False, ".5f"),
|
||
("max_err", "Max err ↓", False, ".5f"),
|
||
("rel_rmse", "Rel-RMSE ↓", False, ".5f"),
|
||
("snr_db", "SNR dB ↑", True, ".2f"),
|
||
("cos_sim", "Cos-sim ↑", True, ".6f"),
|
||
("var_ratio", "Var ratio →1", None, ".5f"),
|
||
("outlier_frac", "Outlier% ↓", False, ".5f"),
|
||
("ch_mse_max", "Ch-MSE max ↓", False, ".5f"),
|
||
("ch_mse_std", "Ch-MSE std ↓", False, ".5f"),
|
||
("delta_mse", "ΔMSE/step ↓", False, ".6f"),
|
||
("delta_cos_sim","ΔCos/step ↑", True, ".7f"),
|
||
]
|
||
|
||
methods = list(all_quant_data.keys())
|
||
if not methods:
|
||
return
|
||
|
||
# Compute stats (same as print_comparison_table)
|
||
stats: dict[str, dict[str, tuple]] = {}
|
||
for method, rows in all_quant_data.items():
|
||
stats[method] = {}
|
||
for col, *_ in TABLE_METRICS:
|
||
if stratify_std:
|
||
stats[method][col] = _stratified_stats(rows, col)
|
||
else:
|
||
stats[method][col] = _nanstats([r[col] for r in rows])
|
||
|
||
std_note = (
|
||
r"± = avg of per-timestep SE (std/√n\_seeds) `[--stratify-std]`"
|
||
if stratify_std else
|
||
r"± = std pooled over all seeds × steps × cond\_types"
|
||
)
|
||
|
||
lines: list[str] = []
|
||
lines.append(f"## Condensed Comparison Table (vs baseline: {baseline_name})\n")
|
||
|
||
# Header row
|
||
header_cells = ["Metric"] + methods
|
||
lines.append("| " + " | ".join(header_cells) + " |")
|
||
|
||
# Alignment row: left for metric label, right-align numeric columns
|
||
align_cells = [":---"] + ["---:" for _ in methods]
|
||
lines.append("| " + " | ".join(align_cells) + " |")
|
||
|
||
# Data rows
|
||
for col, label, higher_is_better, fmt in TABLE_METRICS:
|
||
means = [stats[m][col][0] for m in methods]
|
||
|
||
if higher_is_better is None:
|
||
clean = [(i, v) for i, v in enumerate(means) if v == v]
|
||
winners = {min(clean, key=lambda x: abs(x[1] - 1.0))[0]} if clean else set()
|
||
else:
|
||
winners = _winner_indices(means, higher_is_better)
|
||
|
||
cells = [label]
|
||
for i, method in enumerate(methods):
|
||
mean, std, n = stats[method][col]
|
||
if n == 0:
|
||
cells.append("N/A")
|
||
else:
|
||
star = " ★" if i in winners else ""
|
||
cells.append(f"`{mean:{fmt}} ±{std:{fmt}}`{star}")
|
||
|
||
lines.append("| " + " | ".join(cells) + " |")
|
||
|
||
lines.append("")
|
||
lines.append(f"> ★ = best value for that metric | {std_note}")
|
||
lines.append("")
|
||
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(lines) + "\n")
|
||
|
||
print(f"\nMarkdown table written to: {out_path}")
|
||
|
||
|
||
# ── main ──────────────────────────────────────────────────────────────────────
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Analyze quantization degradation across diffusion timesteps."
|
||
)
|
||
parser.add_argument("--eval-dir", required=True, help="Master folder with one subfolder per precision.")
|
||
parser.add_argument("--baseline", default="BF16", help="Name of the baseline subfolder.")
|
||
parser.add_argument("--out-csv", default="analysis_results.csv", help="Output CSV filename (inside --eval-dir).")
|
||
parser.add_argument("--out-md", default="comparison_table.md", help="Output GitHub Markdown table filename (inside --eval-dir). Pass empty string to skip.")
|
||
parser.add_argument("--stratify-std", action="store_true",
|
||
help="Report SE (std/√n_seeds) across seeds at each fixed timestep "
|
||
"(then average those SEs), rather than pooling all seeds × steps. "
|
||
"Much tighter bars for scale-variant metrics like MSE.")
|
||
args = parser.parse_args()
|
||
|
||
eval_dir = args.eval_dir
|
||
baseline_dir = os.path.join(eval_dir, args.baseline)
|
||
|
||
for path, label in [(eval_dir, "Evaluation directory"), (baseline_dir, "Baseline directory")]:
|
||
if not os.path.exists(path):
|
||
print(f"Error: {label} '{path}' does not exist.")
|
||
return
|
||
|
||
_SIZE_RE = re.compile(r"_?(\d+(?:\.\d+)?)\s*[Gg][Bb]$")
|
||
|
||
raw_subdirs = sorted(
|
||
d for d in os.listdir(eval_dir)
|
||
if os.path.isdir(os.path.join(eval_dir, d)) and d != args.baseline
|
||
)
|
||
|
||
if not raw_subdirs:
|
||
print(f"No quantization folders found to compare against '{args.baseline}'.")
|
||
return
|
||
|
||
# Map display_name → (actual_folder, file_size_gb | None)
|
||
# Size suffix is stripped from the display name used everywhere in output.
|
||
dir_map: dict[str, str] = {} # display_name → actual folder name
|
||
dir_sizes: dict[str, float | None] = {} # display_name → GB (or None)
|
||
|
||
for d in raw_subdirs:
|
||
m = _SIZE_RE.search(d)
|
||
if m:
|
||
display = d[: m.start()] # strip "_8.3GB" suffix
|
||
size_gb = float(m.group(1))
|
||
else:
|
||
display = d
|
||
size_gb = None
|
||
dir_map[display] = d
|
||
dir_sizes[display] = size_gb
|
||
|
||
subdirs = list(dir_map.keys()) # display names, used everywhere below
|
||
|
||
# ── index baseline files ──────────────────────────────────────────────────
|
||
# baseline_grouped[(seed, is_uncond)][step] = filepath
|
||
|
||
baseline_grouped: dict[tuple, dict[int, str]] = defaultdict(dict)
|
||
skipped = 0
|
||
|
||
for filepath in glob.glob(os.path.join(baseline_dir, "*.safetensors")):
|
||
parsed = parse_filename(os.path.basename(filepath))
|
||
if parsed is None:
|
||
print(f" [Skip] Unrecognised filename: {os.path.basename(filepath)}")
|
||
skipped += 1
|
||
continue
|
||
seed, step, is_uncond = parsed
|
||
baseline_grouped[(seed, is_uncond)][step] = filepath
|
||
|
||
print(f"Indexed {sum(len(v) for v in baseline_grouped.values())} baseline files "
|
||
f"({skipped} skipped) across {len(baseline_grouped)} (seed, cond_type) chains.")
|
||
print(f"Comparing: {', '.join(subdirs)}\n")
|
||
|
||
# ── CSV header ───────────────────────────────────────────────────────────
|
||
METRIC_COLS = [
|
||
"mse", "mae", "max_err", "rel_rmse", "snr_db",
|
||
"cos_sim",
|
||
"ch_mse_mean", "ch_mse_std", "ch_mse_max", "worst_channel",
|
||
"var_ratio", "outlier_frac",
|
||
# marginal deltas
|
||
"delta_mse", "delta_mae", "delta_snr_db", "delta_cos_sim", "delta_outlier_frac",
|
||
]
|
||
csv_rows = [["quant_method", "seed", "step", "cond_type"] + METRIC_COLS]
|
||
|
||
# ── process each quantization method ─────────────────────────────────────
|
||
all_quant_data: dict[str, list[dict]] = {}
|
||
|
||
for quant_name in subdirs:
|
||
quant_dir = os.path.join(eval_dir, dir_map[quant_name])
|
||
rows = []
|
||
matches = 0
|
||
misses = 0
|
||
|
||
print(f"--- Processing {quant_name} ---")
|
||
|
||
for (seed, is_uncond), steps_dict in baseline_grouped.items():
|
||
cond_type = "uncond" if is_uncond else "cond"
|
||
sorted_steps = sorted(steps_dict.keys())
|
||
|
||
prev_metrics: dict | None = None
|
||
prev_step_had_miss = False
|
||
|
||
for step in sorted_steps:
|
||
bl_filepath = steps_dict[step]
|
||
filename = os.path.basename(bl_filepath)
|
||
qt_filepath = os.path.join(quant_dir, filename)
|
||
|
||
if not os.path.exists(qt_filepath):
|
||
misses += 1
|
||
prev_metrics = None
|
||
prev_step_had_miss = True
|
||
continue
|
||
|
||
try:
|
||
bl_latent = load_file(bl_filepath)["latent"].to(torch.float32).squeeze()
|
||
qt_latent = load_file(qt_filepath)["latent"].to(torch.float32).squeeze()
|
||
|
||
if bl_latent.shape != qt_latent.shape:
|
||
print(f" [Warning] Shape mismatch for {filename}: "
|
||
f"{bl_latent.shape} vs {qt_latent.shape}")
|
||
misses += 1
|
||
prev_metrics = None
|
||
prev_step_had_miss = True
|
||
continue
|
||
|
||
metrics = compute_metrics(bl_latent, qt_latent)
|
||
|
||
# marginal deltas: only valid when prev step was not missed
|
||
if prev_metrics is not None and not prev_step_had_miss:
|
||
marginal = compute_marginal(metrics, prev_metrics)
|
||
else:
|
||
# mark as NaN so aggregation can ignore them cleanly
|
||
marginal = {k: float("nan") for k in
|
||
("delta_mse", "delta_mae", "delta_snr_db",
|
||
"delta_cos_sim", "delta_outlier_frac")}
|
||
|
||
prev_metrics = metrics
|
||
prev_step_had_miss = False
|
||
|
||
row = {
|
||
"quant_method": quant_name,
|
||
"seed": seed,
|
||
"step": step,
|
||
"cond_type": cond_type,
|
||
**metrics,
|
||
**marginal,
|
||
}
|
||
rows.append(row)
|
||
csv_rows.append([row[c] for c in ["quant_method", "seed", "step", "cond_type"] + METRIC_COLS])
|
||
matches += 1
|
||
|
||
except Exception as e:
|
||
print(f" [Error] Failed processing {filename}: {e}")
|
||
misses += 1
|
||
prev_metrics = None
|
||
prev_step_had_miss = True
|
||
|
||
print(f" Matched {matches} files (Missed {misses}).")
|
||
all_quant_data[quant_name] = rows
|
||
|
||
# ── summary ──────────────────────────────────────────────────────────────
|
||
print("\n" + "═" * 90)
|
||
print("SUMMARY — averaged over all seeds, stratified by step and cond_type")
|
||
print("═" * 90)
|
||
|
||
SUMMARY_METRICS = [
|
||
("mse", "MSE (accum)"),
|
||
("snr_db", "SNR (dB) "),
|
||
("cos_sim", "Cos-sim "),
|
||
("rel_rmse", "Rel-RMSE "),
|
||
("outlier_frac", "Outlier % "),
|
||
("var_ratio", "Var ratio "),
|
||
("delta_mse", "ΔMSE/step "),
|
||
]
|
||
|
||
for quant_name, rows in all_quant_data.items():
|
||
print(f"\n Method: {quant_name}")
|
||
|
||
# group by (step, cond_type)
|
||
by_step: dict[tuple, list[dict]] = defaultdict(list)
|
||
for r in rows:
|
||
by_step[(r["step"], r["cond_type"])].append(r)
|
||
|
||
# per-cond_type global summary
|
||
by_cond: dict[str, list[dict]] = defaultdict(list)
|
||
for r in rows:
|
||
by_cond[r["cond_type"]].append(r)
|
||
|
||
for cond_type, cond_rows in sorted(by_cond.items()):
|
||
print(f"\n [{cond_type}]")
|
||
header = f" {'Metric':<14}" + " Global Avg"
|
||
print(header)
|
||
print(" " + "-" * 30)
|
||
for col, label in SUMMARY_METRICS:
|
||
vals = [r[col] for r in cond_rows if r[col] == r[col]] # skip NaN
|
||
if vals:
|
||
avg = sum(vals) / len(vals)
|
||
print(f" {label} {avg:+.6f}")
|
||
|
||
# worst timesteps by MSE
|
||
step_avg_mse = {
|
||
(step, ct): sum(r["mse"] for r in step_rows) / len(step_rows)
|
||
for (step, ct), step_rows in by_step.items()
|
||
}
|
||
worst = sorted(step_avg_mse.items(), key=lambda x: x[1], reverse=True)[:5]
|
||
print(f"\n Top-5 worst steps by MSE:")
|
||
for (step, ct), val in worst:
|
||
print(f" step={step:>4d} {ct:<6} mse={val:.6f}")
|
||
|
||
# ── condensed comparison table ────────────────────────────────────────────
|
||
print_comparison_table(all_quant_data, args.baseline, stratify_std=args.stratify_std)
|
||
|
||
# ── MSE-vs-bits graph ─────────────────────────────────────────────────────
|
||
plot_path = os.path.join(eval_dir, "mse_vs_bits.png")
|
||
plot_mse_vs_bits(all_quant_data, dir_sizes, plot_path, stratify_std=args.stratify_std)
|
||
|
||
# ── write markdown table ──────────────────────────────────────────────────
|
||
if args.out_md:
|
||
md_path = os.path.join(eval_dir, args.out_md)
|
||
save_markdown_table(all_quant_data, args.baseline, md_path, stratify_std=args.stratify_std)
|
||
|
||
# ── write CSV ─────────────────────────────────────────────────────────────
|
||
if args.out_csv:
|
||
csv_path = os.path.join(eval_dir, args.out_csv)
|
||
with open(csv_path, "w", newline="") as f:
|
||
csv.writer(f).writerows(csv_rows)
|
||
print(f"\nDetailed per-step data written to: {csv_path}")
|
||
|
||
|
||
def plot_mse_vs_bits(
|
||
all_quant_data: dict[str, list[dict]],
|
||
dir_sizes: dict[str, float | None],
|
||
out_path: str,
|
||
stratify_std: bool = False,
|
||
) -> None:
|
||
"""
|
||
Scatter + line plot of mean MSE vs file size (GB) extracted from folder names.
|
||
|
||
Grouping rules
|
||
──────────────
|
||
The prefix before the first digit in the display name is used as the family,
|
||
e.g. "GGUF_Q8_TE" and "GGUF_Q6KM" both belong to "GGUF_Q".
|
||
Dots in the same family share a colour and are connected by a dashed line.
|
||
|
||
Folders with no size suffix are plotted as isolated dots with no line.
|
||
"""
|
||
if not _MATPLOTLIB_OK:
|
||
print("[plot] matplotlib not found – skipping MSE-vs-bits graph. "
|
||
"Install with: pip install matplotlib")
|
||
return
|
||
|
||
BIT_RE = re.compile(r"(\d+)")
|
||
|
||
points: list[tuple[str, float, float, float, str]] = []
|
||
# (display_name, size_gb, mean_mse, std_mse, group)
|
||
|
||
for method, rows in all_quant_data.items():
|
||
size_gb = dir_sizes.get(method)
|
||
if size_gb is None:
|
||
continue # skip methods with no size info
|
||
|
||
m = BIT_RE.search(method)
|
||
prefix = method[: m.start()] if m else method
|
||
group = prefix.rstrip("_-+. ")
|
||
|
||
mean, std, n = (_stratified_stats(rows, "mse") if stratify_std
|
||
else _nanstats([r["mse"] for r in rows]))
|
||
if n == 0:
|
||
continue
|
||
points.append((method, size_gb, mean, std, group))
|
||
|
||
if not points:
|
||
print("[plot] No plottable data found (no GB size suffix in any folder name).")
|
||
return
|
||
|
||
# ── assign one colour per group ───────────────────────────────────────────
|
||
groups = list(dict.fromkeys(p[4] for p in points))
|
||
PALETTE = [
|
||
"#4FC3F7", "#81C784", "#FFB74D", "#F06292",
|
||
"#BA68C8", "#4DB6AC", "#FFF176", "#FF8A65",
|
||
]
|
||
group_colour = {g: PALETTE[i % len(PALETTE)] for i, g in enumerate(groups)}
|
||
|
||
# jitter within same group at same x so dots don't fuse
|
||
from collections import Counter
|
||
seen: Counter = Counter()
|
||
JITTER = 0.05
|
||
|
||
# ── figure ────────────────────────────────────────────────────────────────
|
||
fig, ax = plt.subplots(figsize=(10, 6))
|
||
fig.patch.set_facecolor("#1a1a2e")
|
||
ax.set_facecolor("#16213e")
|
||
for spine in ax.spines.values():
|
||
spine.set_edgecolor("#3a3a5c")
|
||
ax.tick_params(colors="#aaaacc", labelsize=9)
|
||
ax.xaxis.label.set_color("#ccccee")
|
||
ax.yaxis.label.set_color("#ccccee")
|
||
ax.title.set_color("#e0e0ff")
|
||
ax.grid(True, color="#2a2a4a", linewidth=0.7, linestyle="--", alpha=0.8)
|
||
ax.set_axisbelow(True)
|
||
|
||
# lines first (behind dots)
|
||
for group in groups:
|
||
gp = sorted([p for p in points if p[4] == group], key=lambda x: x[1])
|
||
if len(gp) < 2:
|
||
continue
|
||
ax.plot([p[1] for p in gp], [p[2] for p in gp],
|
||
color=group_colour[group], linewidth=1.5,
|
||
linestyle="--", alpha=0.45, zorder=1)
|
||
|
||
# dots + error bars + labels
|
||
label_added = set()
|
||
for method, size_gb, mean, std, group in points:
|
||
colour = group_colour[group]
|
||
key = (group, size_gb)
|
||
x = size_gb + seen[key] * JITTER
|
||
seen[key] += 1
|
||
|
||
lbl = group if group not in label_added else None
|
||
label_added.add(group)
|
||
|
||
ax.errorbar(x, mean, yerr=std,
|
||
fmt="o", color=colour,
|
||
ecolor=colour, elinewidth=1.2, capsize=4, capthick=1.2,
|
||
markersize=8, markeredgewidth=1.2,
|
||
markeredgecolor="#ffffff44",
|
||
label=lbl, zorder=3)
|
||
|
||
ax.annotate(
|
||
method,
|
||
xy=(x, mean), xytext=(0, 10),
|
||
textcoords="offset points",
|
||
ha="center", va="bottom",
|
||
fontsize=7.5, color=colour, alpha=0.92, zorder=4,
|
||
)
|
||
|
||
# axes
|
||
ax.set_xlabel("File size (GB)", labelpad=8)
|
||
ax.set_ylabel("Mean MSE (lower = better)", labelpad=8)
|
||
ax.set_title("MSE vs File Size", pad=14, fontsize=13, fontweight="bold")
|
||
ax.yaxis.set_major_formatter(ticker.FormatStrFormatter("%.5f"))
|
||
|
||
# x ticks at actual sizes
|
||
all_sizes = sorted({p[1] for p in points})
|
||
ax.set_xticks(all_sizes)
|
||
ax.set_xticklabels([f"{s:g} GB" for s in all_sizes])
|
||
|
||
legend = ax.legend(
|
||
title="Family", title_fontsize=9,
|
||
fontsize=8.5, framealpha=0.25,
|
||
facecolor="#1a1a2e", edgecolor="#3a3a5c",
|
||
labelcolor="#ccccee",
|
||
)
|
||
legend.get_title().set_color("#aaaacc")
|
||
|
||
fig.tight_layout(pad=1.6)
|
||
fig.savefig(out_path, dpi=150, bbox_inches="tight",
|
||
facecolor=fig.get_facecolor())
|
||
plt.close(fig)
|
||
print(f"\n[plot] MSE-vs-file-size graph saved to: {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|