| """Export a training checkpoint to the release format: model.safetensors + config.json.""" |
|
|
| |
| |
| |
| |
| import hashlib |
| import json |
| import os |
| import sys |
| from dataclasses import fields |
|
|
| import torch |
| from safetensors.torch import save_file |
|
|
| from facepoints.model import CONFIG, FORMAT, LOG_SCALE_RANGE, WEIGHTS, std_factor |
| from facepoints.train.config import TrainConfig |
|
|
|
|
| def export_release(checkpoint_path, out_dir): |
| """Write model.safetensors and config.json to out_dir from a checkpoint; return the config.""" |
| ck = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| a, classes = ck["args"], [str(c) for c in ck["cls"]] |
| tensors = {f"net.{k}": v.contiguous() for k, v in ck["model"].items()} |
| tensors.update({f"side.{k}": v.contiguous() for k, v in ck.get("side", {}).items()}) |
| tensors.update({f"side_rf.{k}": v.contiguous() for k, v in ck.get("side_rf", {}).items()}) |
| keep = {f.name for f in fields(TrainConfig)} - set(TrainConfig.PATH_KEYS) |
| config = { |
| "format": FORMAT, |
| "backbone": a["backbone"], |
| "input_px": int(a["res"]), |
| "neck": int(a["neck"]), |
| "stride": int(a.get("stride", 8)), |
| "n_points": len(classes), |
| "classes": classes, |
| "likelihood": a["nll"], |
| "nu": float(a.get("nu", 10.0)), |
| "std_factor": std_factor(a["nll"], float(a.get("nu", 10.0))), |
| "log_scale_range": list(LOG_SCALE_RANGE), |
| "side_heads": {k: int(v) for k, v in ck.get("side_k", {}).items()}, |
| "side_refine": "side_rf" in ck, |
| "iterations": int(ck["it"]), |
| "training": {k: v for k, v in a.items() if k in keep}, |
| } |
| os.makedirs(out_dir, exist_ok=True) |
| save_file(tensors, os.path.join(out_dir, WEIGHTS), metadata={"format": "pt"}) |
| with open(os.path.join(out_dir, CONFIG), "w") as fh: |
| json.dump(config, fh, indent=1) |
| return config |
|
|
|
|
| def sha256(path): |
| """Return the SHA-256 hex digest of a file.""" |
| h = hashlib.sha256() |
| with open(path, "rb") as fh: |
| for chunk in iter(lambda: fh.read(1 << 20), b""): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def main(): |
| """Export a checkpoint and print the weights file's size and SHA-256.""" |
| if len(sys.argv) != 3: |
| sys.exit("usage: python -m facepoints.train.export <ckpt.pt> <out dir>") |
| config = export_release(sys.argv[1], sys.argv[2]) |
| w = os.path.join(sys.argv[2], WEIGHTS) |
| mb = os.path.getsize(w) / 1e6 |
| print(f"wrote {w} ({mb:.1f} MB, sha256 {sha256(w)}) and {CONFIG}: {config['n_points']} points, " |
| f"side heads {config['side_heads']}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
