File size: 5,970 Bytes
6fbb543 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Measure exactly how PixelGen changed after adding one class.
Compares the frozen before-state (experiments/before.pt, before_outputs.npz,
before_meta.json) against the live model built from current source.
Reports:
1. tensor shapes before -> after
2. raw parameter diff (how many values differ by index)
3. semantic diff (how many values differ once rows/cols are aligned by
class name / word, i.e. the change that actually matters)
4. vocab index remap (sorted vocab means new words shift old columns)
5. output equivalence on every pre-existing prompt (bit-exact or not)
6. what the new prompts do now vs before
"""
import json
from pathlib import Path
import numpy as np
import torch
from pixelgen.model import PixelGen
HERE = Path(__file__).resolve().parent
before_sd = torch.load(HERE / "before.pt", map_location="cpu", weights_only=True)
before_out = np.load(HERE / "before_outputs.npz")
meta = json.load(open(HERE / "before_meta.json"))
b_vocab, b_classes, prompts = meta["vocab"], meta["classes"], meta["prompts"]
m = PixelGen()
after_sd = m.state_dict()
a_vocab, a_classes = m.vocab, m.classes
new_classes = [c for c in a_classes if c not in b_classes]
new_words = [w for w in a_vocab if w not in b_vocab]
print("=" * 70)
print(f"new classes: {new_classes} new words: {new_words}")
print(f"vocab {len(b_vocab)} -> {len(a_vocab)} classes {len(b_classes)} -> {len(a_classes)}")
# 1. shapes -------------------------------------------------------------
print("\n[1] tensor shapes")
tot_b = tot_a = 0
for k in after_sd:
sb, sa = tuple(before_sd[k].shape), tuple(after_sd[k].shape)
tot_b += before_sd[k].numel(); tot_a += after_sd[k].numel()
flag = " <-- grew" if sb != sa else ""
print(f" {k:24} {str(sb):>16} -> {str(sa):<16}{flag}")
print(f" total values {tot_b} -> {tot_a} (+{tot_a - tot_b})")
# 2. raw diff ------------------------------------------------------------
print("\n[2] raw diff by index (naive view)")
for k in after_sd:
b, a = before_sd[k], after_sd[k]
if b.shape == a.shape:
n = int((b != a).sum())
print(f" {k:24} {n:5d} / {a.numel()} values differ")
else:
# compare overlapping prefix region
sl = tuple(slice(0, s) for s in b.shape)
n = int((b != a[sl]).sum())
print(f" {k:24} {n:5d} / {b.numel()} values differ in the old-shaped prefix (+{a.numel()-b.numel()} new)")
# 3. semantic diff (align by names) ---------------------------------------
print("\n[3] semantic diff (rows/cols aligned by class name / word)")
def aligned_linear(name):
b, a = before_sd[f"{name}.weight"], after_sd[f"{name}.weight"]
rows = b_classes if name == "word_to_class" else m.ramps
changed = []
for ri, r in enumerate(rows):
for w in b_vocab:
bv = b[ri, b_vocab[w]].item()
av = a[a_classes.index(r) if name == "word_to_class" else ri, a_vocab[w]].item()
if bv != av:
changed.append((r, w, bv, av))
return changed
for name in ("word_to_class", "word_to_ramp"):
ch = aligned_linear(name)
print(f" {name:16} old (row,word) entries changed: {len(ch)} {ch[:5]}")
bb, ab = before_sd["word_to_class.bias"], after_sd["word_to_class.bias"]
print(" word_to_class.bias old-class entries changed:",
sum(bb[i].item() != ab[a_classes.index(c)].item() for i, c in enumerate(b_classes)))
for bank in ("on_bank", "free_bank"):
b, a = before_sd[bank], after_sd[bank]
n = sum(int((b[i] != a[a_classes.index(c)]).sum()) for i, c in enumerate(b_classes))
print(f" {bank:16} old-class template cells changed: {n}")
for k in ("shade_bank", "ramp_table", "outline.weight", "palette_conv.weight"):
print(f" {k:16} identical: {torch.equal(before_sd[k], after_sd[k])}")
# 4. vocab remap ----------------------------------------------------------
print("\n[4] vocab index remap caused by sorted insertion")
moved = [(w, b_vocab[w], a_vocab[w]) for w in b_vocab if b_vocab[w] != a_vocab[w]]
print(f" {len(moved)} of {len(b_vocab)} old words changed column index")
print(" first few:", moved[:6])
print(" new word columns:", {w: a_vocab[w] for w in new_words})
# 5. output equivalence ---------------------------------------------------
print("\n[5] output equivalence on pre-existing prompts (3 seeds each)")
same = diff = 0
changed_prompts = []
for i, p in enumerate(prompts):
b = before_out[f"p{i}"]
a = m.generate([p] * 3, [0, 1, 2], scale=1)
if np.array_equal(a, b):
same += 1
else:
diff += 1
changed_prompts.append(p)
print(f" identical: {same} changed: {diff} -> {changed_prompts}")
# 6. new prompt behaviour ---------------------------------------------------
print("\n[6] what the new-class prompts resolve to now")
ref = {c: (m.generate([c], [0], scale=1)[0].sum(-1) > 0) for c in a_classes}
for p in ["mushroom", "red mushroom", "purple fungus", "toadstool", "dragon", ""]:
mask = m.generate([p], [0], scale=1)[0].sum(-1) > 0
cls = [c for c, r in ref.items() if np.array_equal(r, mask)]
i = prompts.index(p) if p in prompts else None
was = "n/a"
if i is not None:
bm = before_out[f"p{i}"][0].sum(-1) > 0
was = [c for c, r in ref.items() if np.array_equal(r, bm)]
print(f" {p!r:18} before -> {was} after -> {cls}")
# image ---------------------------------------------------------------------
from PIL import Image
S = 6; cell = 16 * S; pad = 4
rows = [("before", b_classes), ("after", a_classes)]
W = max(len(c) for _, c in rows) * (cell + pad) + pad
img = Image.new("RGB", (W, 2 * (cell + pad) + pad), (29, 43, 83))
for r, (_, classes) in enumerate(rows):
for c, cls in enumerate(classes):
img.paste(Image.fromarray(m.generate([cls], [0], scale=S)[0]), (pad + c * (cell + pad), pad + r * (cell + pad)))
img.save(HERE / "before_after.png")
print(f"\nwrote {HERE/'before_after.png'} (top: old classes rendered by new model; bottom: all classes)")
|