| """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}") |
|
|
| |
| 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", |
| ) |
|
|
| |
| 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", |
| ) |
|
|
| |
| 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", |
| ) |
|
|
| |
| 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") |
|
|
