File size: 1,269 Bytes
2847d0b | 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 | #!/usr/bin/env python3
"""Validate a Predictor training_latest.pt completion state."""
from __future__ import annotations
import argparse
from pathlib import Path
import torch
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("path", type=Path)
parser.add_argument("--expected_step", type=int, required=True)
parser.add_argument("--supervision_mode", choices=("onpolicy", "offline_ffff"))
args = parser.parse_args()
if not args.path.is_file():
raise FileNotFoundError(args.path)
state = torch.load(args.path, map_location="cpu", weights_only=False)
actual = int(state.get("global_step", -1))
if actual != args.expected_step:
raise ValueError(f"{args.path}: global_step={actual}, expected {args.expected_step}")
if args.supervision_mode is not None:
saved = str(state.get("supervision_mode"))
if saved != args.supervision_mode:
raise ValueError(
f"{args.path}: supervision_mode={saved!r}, "
f"expected {args.supervision_mode!r}"
)
print(
f"valid training state: path={args.path} step={actual} "
f"mode={state.get('supervision_mode', 'stage1')}"
)
if __name__ == "__main__":
main()
|