| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from common import ( |
| build_model, |
| collect_predictions, |
| get_device, |
| load_config, |
| make_loaders, |
| mean_relative_l2_percent, |
| resolve_output_path, |
| save_json, |
| set_seed, |
| ) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(Path(__file__).resolve().parents[1] / "config" / "config.yaml")) |
| parser.add_argument("--checkpoint", default=None) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| cfg = load_config(args.config) |
| set_seed(cfg["training"]["seed"], cfg["runtime"]["deterministic"], cfg["runtime"]["benchmark"]) |
| device = get_device(cfg) |
| checkpoint = Path(args.checkpoint) if args.checkpoint else resolve_output_path(cfg, cfg["paths"]["checkpoint"]) |
| results_dir = resolve_output_path(cfg, cfg["paths"]["results_dir"]) |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| _, _, test_loader = make_loaders(cfg) |
| model = build_model(cfg).to(device) |
| payload = torch.load(checkpoint, map_location=device) |
| model.load_state_dict(payload["model_state"] if "model_state" in payload else payload) |
| test_l2 = mean_relative_l2_percent(model, test_loader, device) |
| inputs, preds, targets = collect_predictions(model, test_loader, device) |
| pred_path = results_dir / "test_predictions.npz" |
| np.savez_compressed(pred_path, inputs=inputs, predictions=preds, targets=targets) |
| save_json( |
| results_dir / "inference_metrics.json", |
| { |
| "status": "complete", |
| "checkpoint": str(checkpoint), |
| "prediction_file": str(pred_path), |
| "test_mean_relative_l2_percent": test_l2, |
| "paper_crop_re10000_percent": cfg["experiment"]["paper_reference_test_percent"], |
| "paper_crop_re10000_std_percent": cfg["experiment"]["paper_reference_std_percent"], |
| }, |
| ) |
| print(f"inference_complete test_mean_relative_l2={test_l2:.6f}% predictions={pred_path}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
