File size: 2,878 Bytes
f930dac | 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 | """Load a published checkpoint.
Each published run is a directory holding
model.safetensors the parameters
config.json {"kind", "cfg", "train_args", "params", ...}
`kind` selects the architecture; `cfg` is the exact dataclass field set the run
was built with, so nothing is inferred from tensor shapes.
"""
from __future__ import annotations
import json
import torch
from safetensors.torch import load_file
from .common import ckpt_root
def load_config(run: str) -> dict:
return json.loads((ckpt_root() / run / "config.json").read_text())
def build(kind: str, cfg: dict):
if kind == "rmt":
from .model.rmt import RMT, RMTConfig
return RMT(RMTConfig(**{k: v for k, v in cfg.items()
if k in RMTConfig.__dataclass_fields__}))
if kind in ("pnsr", "pnsr_k1"):
from .model.pnsr import PNSR, PNSRConfig
return PNSR(PNSRConfig(**{k: v for k, v in cfg.items()
if k in PNSRConfig.__dataclass_fields__}))
if kind in ("bind", "unbound"):
from .model.bind import BindConfig, PNSBind
return PNSBind(BindConfig(**{k: v for k, v in cfg.items()
if k in BindConfig.__dataclass_fields__}))
from .model.tx import TX, TXConfig
return TX(TXConfig(**{k: v for k, v in cfg.items()
if k in TXConfig.__dataclass_fields__}))
def load_model(run: str, device: str = "cuda"):
"""Returns (model in eval mode on `device`, kind, config dict)."""
conf = load_config(run)
model = build(conf["kind"], conf["cfg"])
state = load_file(str(ckpt_root() / run / "model.safetensors"), device="cpu")
model.load_state_dict(state, strict=True)
model.eval().to(device)
return model, conf["kind"], conf
def param_count(model) -> int:
"""Number of distinct trainable parameters — the paper's budget figure.
The event encoder and the record encoder *share* one 8192x384 token
embedding table (the parameter is tied, not duplicated), so `state_dict`
lists it under two keys while the model holds it once. `parameters()`
already de-duplicates shared tensors, which is exactly the quantity the
matched-parameter comparison uses.
"""
return sum(p.numel() for p in model.parameters())
@torch.no_grad()
def verify(run: str, device: str = "cpu") -> dict:
"""Load a checkpoint and check it against its recorded provenance."""
conf = load_config(run)
model, _, _ = load_model(run, device)
n = param_count(model)
tensors = sum(v.numel() for v in model.state_dict().values())
return dict(run=run, kind=conf["kind"], loaded=True,
params=n, params_match=n == conf["params"],
state_dict_tensor_total=tensors,
tensor_total_match=tensors == conf["state_dict_tensor_total"])
|