61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
"""Reading a published weight release back into a model.
|
|
|
|
A release holds the trained weights alone: no optimizer state, and no copy of
|
|
the frozen frontend, which is restored from its own pretrained source when the
|
|
model is built. It also ships the model half of the config it was trained with,
|
|
so a caller supplies only where the data lives and where results go.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import yaml
|
|
|
|
WEIGHTS_NAME = "weights.pt"
|
|
CONFIG_NAME = "config.yaml"
|
|
PROVENANCE_NAME = "provenance.json"
|
|
|
|
|
|
def resolve(path):
|
|
"""Accept either a release directory or the weights file inside one."""
|
|
path = Path(path)
|
|
if path.is_dir():
|
|
weights = path / WEIGHTS_NAME
|
|
if not weights.exists():
|
|
raise SystemExit(f"{path} is not a release: no {WEIGHTS_NAME}")
|
|
return weights, path / CONFIG_NAME
|
|
return path, path.parent / CONFIG_NAME
|
|
|
|
|
|
def merge_config(cfg, release_config, module):
|
|
"""Take the architecture from the release, keep local paths from `cfg`.
|
|
|
|
A release must be rebuilt with the architecture it was trained with, but
|
|
dataset roots and output directories belong to whoever is running it.
|
|
"""
|
|
if not Path(release_config).exists():
|
|
return cfg
|
|
with open(release_config) as handle:
|
|
published = yaml.safe_load(handle) or {}
|
|
if module in published:
|
|
cfg[module] = {**cfg.get(module, {}), **published[module]}
|
|
return cfg
|
|
|
|
|
|
def load_weights(model, weights_path):
|
|
"""Load a release into `model`, allowing only the frozen frontend to be absent."""
|
|
state = torch.load(weights_path, map_location="cpu", weights_only=True)
|
|
missing, unexpected = model.load_state_dict(state, strict=False)
|
|
if unexpected:
|
|
raise SystemExit(
|
|
f"{weights_path} holds weights this model does not want, so it was trained "
|
|
f"with a different architecture: {sorted(unexpected)[:5]}"
|
|
)
|
|
frozen = {name for name, p in model.named_parameters() if not p.requires_grad}
|
|
unexplained = set(missing) - frozen
|
|
if unexplained:
|
|
raise SystemExit(
|
|
f"{weights_path} is missing trainable weights: {sorted(unexplained)[:5]}"
|
|
)
|
|
return model
|