File size: 4,630 Bytes
ffdcfe7 | 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 | """Idempotently graft the MAE objective into the eatmap tree on the VM.
Kept here rather than applied by hand because /workspace is not a volume on this
instance -- workspace_is_volume is False, so a recycle wipes the container
filesystem. This script plus mae.py is the whole change; re-running it on a
fresh box reproduces the tree.
Edits are surgical and leave every EAT code path byte-identical: the objective
dispatch defaults to "ufo", and the two structural fixes (teacher-optional
export, teacher-aware parameter count) are no-ops when a teacher is present.
"""
import sys
from pathlib import Path
ROOT = Path(sys.argv[1] if len(sys.argv) > 1 else "/workspace/code/eat-map-regmix/eatmap")
def patch(path: Path, old: str, new: str, label: str) -> None:
text = path.read_text()
if new in text:
print(f" [skip] {label} (already applied)")
return
if old not in text:
raise SystemExit(f" [FAIL] {label}: anchor not found in {path}")
path.write_text(text.replace(old, new, 1))
print(f" [ok] {label}")
print(f"patching {ROOT}")
# ---- 1. ObjectiveConfig: objective selector + MAE decoder geometry ----------
patch(
ROOT / "config.py",
""" # null/0 means "anneal over total_steps" -- never leave this at a value
# larger than the budget or the EMA never reaches its end decay.
ema_anneal_steps: int | None = None""",
""" # null/0 means "anneal over total_steps" -- never leave this at a value
# larger than the budget or the EMA never reaches its end decay.
ema_anneal_steps: int | None = None
# "ufo" is EAT's EMA-teacher objective and the default, so every existing
# config keeps its meaning. "mae" swaps in masked log-mel reconstruction
# over the same encoder -- see eatmap/mae.py for why that sharing matters.
type: str = "ufo"
# Decoder geometry, MAE only. Asymmetric by design: the decoder is thrown
# away after pretraining, so it is narrow and shallow relative to the
# encoder. Defaults are sized for the 384-dim 15M proxy.
mae_decoder_dim: int = 256
mae_decoder_depth: int = 4
mae_decoder_heads: int = 8
# Per-patch target standardization (He et al.'s norm_pix_loss). Off makes
# the loss track patch loudness rather than structure.
mae_norm_pix: bool = True""",
"ObjectiveConfig: type + MAE decoder fields",
)
# ---- 2. export_weights: teacher is optional --------------------------------
patch(
ROOT / "runner.py",
""" tensors = {}
for name, tensor in model.student.state_dict().items():
tensors[f"student.{name}"] = tensor.detach().cpu().contiguous()
for name, tensor in model.teacher.state_dict().items():
tensors[f"teacher.{name}"] = tensor.detach().cpu().contiguous()
for name, tensor in model.decoder.state_dict().items():
tensors[f"decoder.{name}"] = tensor.detach().cpu().contiguous()""",
""" tensors = {}
# MAE has no teacher, so the parts are discovered rather than assumed. The
# student is written under the same "student." prefix either way, which is
# what probe.load_encoder keys on -- an MAE export and an EAT export are
# interchangeable to every readout.
for part in ("student", "teacher", "decoder"):
module = getattr(model, part, None)
if module is None:
continue
for name, tensor in module.state_dict().items():
tensors[f"{part}.{name}"] = tensor.detach().cpu().contiguous()""",
"export_weights: teacher-optional",
)
# ---- 3. objective dispatch -------------------------------------------------
patch(
ROOT / "runner.py",
"from .model import EATPretrainer, parameter_counts",
"from .mae import MAEPretrainer\nfrom .model import EATPretrainer, parameter_counts",
"runner: import MAEPretrainer",
)
patch(
ROOT / "runner.py",
" model = EATPretrainer(config).to(device)",
""" pretrainer = {"ufo": EATPretrainer, "mae": MAEPretrainer}.get(config.objective.type)
if pretrainer is None:
raise ValueError(f"unknown objective.type {config.objective.type!r}")
model = pretrainer(config).to(device)""",
"runner: objective dispatch",
)
# ---- 4. parameter_counts: only EAT carries a frozen twin -------------------
patch(
ROOT / "model.py",
""" "total_params": student * 2 + decoder, # teacher is a frozen copy""",
""" # EAT carries a frozen EMA twin of the student; MAE does not.
"total_params": student * (2 if getattr(model, "teacher", None) is not None else 1) + decoder,""",
"parameter_counts: teacher-aware total",
)
print("done")
|