| """Rename a trained model's choice labels without retraining the policy. |
| |
| The core only sees choice *embeddings*. To move a model from "UP" to "move up" |
| we fit the ChoiceEncoder (not the shared embeddings, not the core) so that the |
| new label produces the vector the old label produced. The decision function is |
| then unchanged up to the fit error, which is verified on real states. |
| |
| minijev relabel --checkpoint runs/x/final.pt --out runs/x/final-semantic.pt \\ |
| --map UP="move up" DOWN="move down" LEFT="move left" RIGHT="move right" \\ |
| --set task.choice_style=semantic |
| """ |
| from __future__ import annotations |
|
|
| import copy |
|
|
| import numpy as np |
| import torch |
|
|
|
|
| def relabel_checkpoint(checkpoint: str, out: str, mapping: dict[str, str], cfg_updates: dict | None = None, |
| steps: int = 3000, lr: float = 3e-3, anchors: list[str] | None = None) -> dict: |
| from minijev.inference.runtime import TorchDecider |
| from minijev.training.supervised import make_model |
|
|
| ck = torch.load(checkpoint, map_location="cpu", weights_only=False) |
| cfg = copy.deepcopy(ck["config"]) |
| model = make_model(cfg) |
| model.load_state_dict(ck["model"]) |
| model.eval() |
| old_dec = TorchDecider(model=copy.deepcopy(model), cfg=cfg) |
| tok, max_len = old_dec.tokenizer, old_dec.max_choice_len |
| old, new = list(mapping), list(mapping.values()) |
| anchors = list(anchors or []) |
|
|
| def emb(labels): |
| if model.backbone is None: |
| ids = torch.as_tensor(tok.encode_batch(labels, max_len)) |
| return model.choice_encoder(ids) |
| ids = torch.as_tensor(tok.encode_batch(labels, max_len, dynamic=True)) |
| return model.choice_encoder(pooled=model.backbone(ids)[2]) |
|
|
| with torch.no_grad(): |
| target = emb(old + anchors).detach() |
| shared = {id(p) for n, p in model.named_parameters() if not n.startswith("choice_encoder.")} |
| params = [p for p in model.choice_encoder.parameters() if id(p) not in shared] |
| for p in model.parameters(): |
| p.requires_grad_(False) |
| for p in params: |
| p.requires_grad_(True) |
| opt = torch.optim.Adam(params, lr=lr) |
| sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, steps) |
| for _ in range(steps): |
| loss = (emb(new + anchors) - target).pow(2).mean() / target.pow(2).mean() |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| sched.step() |
| for p in model.parameters(): |
| p.requires_grad_(True) |
|
|
| for k, v in (cfg_updates or {}).items(): |
| node = cfg |
| parts = k.split(".") |
| for part in parts[:-1]: |
| node = node.setdefault(part, {}) |
| node[parts[-1]] = v |
| new_dec = TorchDecider(model=model, cfg=cfg) |
| new_dec.temperature = old_dec.temperature |
| states = new_dec.task.example_states(256, seed=123) |
| p_old = old_dec.probs(states, new_dec.task.default_question or "", old) |
| p_new = new_dec.probs(states, new_dec.task.default_question or "", new) |
| report = {"rel_mse": float(loss.detach()), "max_abs_prob_diff": float(np.abs(p_old - p_new).max()), |
| "argmax_agreement": float((p_old.argmax(1) == p_new.argmax(1)).mean()), "states": len(states)} |
| torch.save({**{k: v for k, v in ck.items() if k not in ("model", "config", "optimizer")}, |
| "model": model.state_dict(), "config": cfg, "relabel": {"from": checkpoint, "map": mapping, **report}}, |
| out) |
| return report |
|
|
