File size: 3,095 Bytes
534cc76 | 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 | """Save decision-head weights and validate model and calibration identities."""
import hashlib
from pathlib import Path
import torch
from .provenance import RUNTIME_CONTRACT
from .atomic import atomic_output
def learned_state(model):
return {"head": {key: tensor.detach().cpu().clone() for key, tensor in model.head.state_dict().items()}}
def learned_digest(model):
digest = hashlib.sha256()
for module, state in learned_state(model).items():
for name, tensor in sorted(state.items()):
digest.update(f"{module}.{name}:{tensor.dtype}:{tuple(tensor.shape)}".encode())
digest.update(tensor.contiguous().numpy().tobytes())
return digest.hexdigest()
def save_checkpoint(path, model, training=None):
"""Atomically save learned weights and optional resumable training state."""
temperature = float(model.temperature)
if not 0 < temperature < float("inf"):
raise ValueError("Invalid checkpoint temperature")
if model.calibration is None and temperature != 1.0:
raise ValueError("A non-unit temperature needs outcome-fitted calibration provenance")
if model.calibration and model.calibration["learned_sha256"] != learned_digest(model):
raise ValueError("Calibration does not match current decision-head weights")
if model.calibration and abs(model.calibration["temperature"]-temperature) > 1e-7:
raise ValueError("Calibration temperature changed")
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
payload = {"contract": RUNTIME_CONTRACT, "identity": model.identity(),
"learned": learned_state(model), "temperature": temperature,
"calibration": model.calibration, "training": training}
with atomic_output(path) as output:
torch.save(payload, output)
def load_checkpoint(path, model):
payload = torch.load(path, map_location="cpu", weights_only=True)
if payload.get("contract") != RUNTIME_CONTRACT or payload.get("identity") != model.identity():
raise ValueError("Checkpoint is not the exact current architecture/assets/preparation")
if set(payload.get("learned", {})) != {"head"}:
raise ValueError("Checkpoint must contain only decision-head weights")
model.head.load_state_dict(payload["learned"]["head"], strict=True)
temperature = float(payload["temperature"])
if not 0 < temperature < float("inf"):
raise ValueError("Invalid checkpoint temperature")
calibration = payload.get("calibration")
if calibration is None and temperature != 1.0:
raise ValueError("A non-unit temperature needs outcome-fitted calibration provenance")
if calibration and (calibration["learned_sha256"] != learned_digest(model)
or abs(calibration["temperature"]-temperature) > 1e-7):
raise ValueError("Calibration binding differs from the saved weights/temperature")
model.temperature.fill_(temperature)
model.calibration = calibration
model.eval()
return payload
|