mirror of
https://github.com/unanmed/ginka-generator.git
synced 2026-08-14 18:12:28 +08:00
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import sys
|
||
import torch
|
||
|
||
if len(sys.argv) < 2:
|
||
print("用法: python diagnose_codebook.py result/seperated/sep-XXX.pth")
|
||
sys.exit(1)
|
||
|
||
ckpt_path = sys.argv[1]
|
||
ckpt = torch.load(ckpt_path, map_location="cpu")
|
||
|
||
names = ["quantizer1", "quantizer2", "quantizer3"]
|
||
labels = ["q1(stage1墙)", "q2(stage2门/怪/入口)", "q3(stage3资源)"]
|
||
|
||
print(f"Checkpoint: {ckpt_path}")
|
||
print(f"Epoch: {ckpt.get('epoch', '?')}")
|
||
print()
|
||
print(f"{'名称':<20s} {'存活/总量':>10s} {'利用率':>8s} {'perplexity':>12s} {'最近大小中位':>14s}")
|
||
print("-" * 68)
|
||
|
||
for name, label in zip(names, labels):
|
||
sd = ckpt.get(name)
|
||
if sd is None:
|
||
print(f"{label:<20s} {'未找到':>10s}")
|
||
continue
|
||
cs = sd["ema_cluster_size"]
|
||
K = cs.numel()
|
||
alive = int((cs > 1.0).sum())
|
||
p = cs / cs.sum()
|
||
ppl = float(torch.exp(-(p * torch.log(p.clamp_min(1e-10))).sum()))
|
||
usage = alive / K * 100
|
||
|
||
# 最近（cluster_size 中位数，反映码字被使用频率）
|
||
median_size = float(cs.median())
|
||
|
||
print(f"{label:<20s} {alive:>3d}/{K:<3d} {usage:>5.1f}% {ppl:>8.2f}/{K:<8d} {median_size:>10.4f}")
|
||
|
||
# 检查是否有码字范数为零的死码
|
||
print()
|
||
print("--- 死码检查（范数=0 的码字） ---")
|
||
for name, label in zip(names, labels):
|
||
sd = ckpt.get(name)
|
||
if sd is None:
|
||
continue
|
||
w = sd["codebook.weight"]
|
||
if w is None:
|
||
w = sd.get("weight") # 可能的备用 key
|
||
if w is None:
|
||
print(f"{label}: 无法读取 codebook.weight")
|
||
continue
|
||
norms = w.norm(dim=1)
|
||
dead = int((norms < 1e-6).sum().item())
|
||
if dead > 0:
|
||
print(f"{label}: {dead} / {w.size(0)} 死码（范数=0）")
|
||
else:
|
||
print(f"{label}: 全部码字有非零范数")
|