| |
| """Strip a training checkpoint down to what inference needs. |
| |
| A training checkpoint carries the optimizer and AMP-scaler state so a run can |
| resume: 795 MB, of which 297 MB is the model. Inference never reads the rest, |
| and for a hosted demo the difference is the whole cold start. |
| |
| python scripts/export_inference_weights.py \ |
| --checkpoint weights/segmentation/fpw_aff_base_ft_512_slits_pgbmi.pth \ |
| --output dist/weights/fpw_aff_base_ft_512_slits_pgbmi.pth |
| |
| Keeps ``model_state_dict`` plus the small scalar metadata (epoch, test_miou), |
| so the result still loads through the same path as the original. |
| """ |
|
|
| import argparse |
| import logging |
| import os |
| import sys |
|
|
| import torch |
|
|
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| |
| KEEP = ("model_state_dict", "epoch", "test_miou") |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def strip_checkpoint(payload: dict) -> dict: |
| """Return a copy of ``payload`` with only the inference-relevant keys. |
| |
| Args: |
| payload: a loaded checkpoint dict. |
| Returns: |
| A new dict holding the keys in :data:`KEEP` that were present. |
| Raises: |
| KeyError: if there is no model state to keep, which would make the |
| output silently useless. |
| """ |
| if "model_state_dict" not in payload: |
| raise KeyError( |
| f"no model_state_dict in this checkpoint; keys are " |
| f"{sorted(payload)}. Nothing would be left to load.") |
| return {key: payload[key] for key in KEEP if key in payload} |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser( |
| description=__doc__, |
| formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--checkpoint", required=True, |
| help="Training checkpoint to read.") |
| parser.add_argument("--output", required=True, |
| help="Where to write the stripped copy.") |
| args = parser.parse_args() |
|
|
| logging.basicConfig(level=logging.INFO, format="%(message)s") |
|
|
| from affmae.eval.loader import legacy_checkpoint_compat |
|
|
| with legacy_checkpoint_compat(): |
| payload = torch.load(args.checkpoint, map_location="cpu", |
| weights_only=False) |
| if not isinstance(payload, dict): |
| raise SystemExit( |
| f"{args.checkpoint} holds a {type(payload).__name__}, not a dict; " |
| f"it is already a bare state dict and needs no stripping.") |
|
|
| kept = strip_checkpoint(payload) |
| dropped = sorted(set(payload) - set(kept)) |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(args.output)) or ".", |
| exist_ok=True) |
| torch.save(kept, args.output) |
|
|
| before = os.path.getsize(args.checkpoint) / 1048576 |
| after = os.path.getsize(args.output) / 1048576 |
| logger.info("%s -> %s", args.checkpoint, args.output) |
| logger.info(" %.1f MB -> %.1f MB (%.1fx smaller), dropped %s", |
| before, after, before / after if after else 0, dropped) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
