Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192#!/usr/bin/env python3"""Evaluate a saved checkpoint on a dataset — standalone, no training.
Used for the authoritative full-test-set number at the end of a run (in-trainingevals use a subsample for speed), and as the shared measuring instrument forcomparing DSEM colonies against TRM baselines.
Usage: python scripts/eval_checkpoint.py \ --checkpoint checkpoints/Sudoku-.../repro-mlp-3/step_65104 \ --config checkpoints/Sudoku-.../repro-mlp-3/all_config.yaml \ --data data/sudoku-extreme-1k-aug-1000"""
import argparseimport jsonimport osimport sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torchimport yaml
from dsem.device import get_devicefrom dsem.pretrain import PretrainConfig, create_dataloader, create_model
def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--checkpoint", required=True) ap.add_argument("--config", required=True, help="all_config.yaml saved with the run") ap.add_argument("--data", required=True, help="dataset dir containing test/") ap.add_argument("--batch-size", type=int, default=None, help="override global_batch_size") ap.add_argument("--out", default=None, help="write metrics JSON here") args = ap.parse_args()
device = get_device()
with open(args.config) as f: cfg_dict = yaml.unsafe_load(f) config = PretrainConfig(**cfg_dict) config.data_paths_test = [args.data] config.evaluators = [] if args.batch_size: config.global_batch_size = args.batch_size
eval_loader, eval_metadata = create_dataloader( config, "test", test_set_mode=True, epochs_per_iter=1, global_batch_size=config.global_batch_size, rank=0, world_size=1, )
model, _, _ = create_model(config, eval_metadata, rank=0, world_size=1) state = torch.load(args.checkpoint, map_location=device) model.load_state_dict(state, assign=True) model.eval()
totals: dict[str, torch.Tensor] = {} batches = 0 with torch.inference_mode(): for _set_name, batch, _gbs in eval_loader: batch = {k: v.to(device) for k, v in batch.items()} with torch.device(device): carry = model.initial_carry(batch) while True: carry, _loss, metrics, _preds, all_finish = model(carry=carry, batch=batch, return_keys=[]) if all_finish: break for k, v in metrics.items(): totals[k] = totals.get(k, torch.zeros((), device=device, dtype=torch.float32)) + v.to(torch.float32) batches += 1 if batches % 25 == 0: seen = int(totals["count"].item()) print(f" {batches} batches, {seen} puzzles, " f"running exact-acc {totals['exact_accuracy'].item()/max(seen,1)*100:.2f}%", flush=True)
count = max(totals.pop("count").item(), 1) results = {k: v.item() / count for k, v in totals.items()} results["puzzles"] = count results["checkpoint"] = args.checkpoint results["data"] = args.data
print(json.dumps(results, indent=2)) if args.out: os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) with open(args.out, "w") as f: json.dump(results, f, indent=2)
if __name__ == "__main__": main()