export_warm_init.py

"""Export a weights-only init from a training bundle: live state overlaid with the
EMA shadows (the eval-time weights), _orig_mod stripped — the format load_checkpoint
expects for warm-starting a new run."""
import sys, torch

bundle, out = sys.argv[1], sys.argv[2]
d = torch.load(bundle, map_location="cpu", weights_only=False)
strip = lambda sd: {k.replace("_orig_mod.", ""): v for k, v in sd.items()}
sd = strip(d["model"])
sd.update(strip(d["ema"]))
torch.save(sd, out)
print(f"{out}: {len(sd)} keys from {bundle}")