File size: 2,157 Bytes
8bf8ea0 | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | 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()
|