| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from common import ( |
| DEFAULT_CONFIG, |
| build_datapipe, |
| build_model, |
| cleanup_distributed, |
| get_attr, |
| initialize_distributed, |
| load_config, |
| load_model_state, |
| predict_batch, |
| prepare_config, |
| ) |
|
|
|
|
| @torch.no_grad() |
| def main() -> int: |
| parser = argparse.ArgumentParser(description="Run PDENNEval FNO inference.") |
| parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml") |
| parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir") |
| parser.add_argument("--checkpoint", default=None, help="Checkpoint path") |
| parser.add_argument("--output-dir", default=None, help="Directory for npz predictions") |
| parser.add_argument("--max-batches", type=int, default=None) |
| parser.add_argument("--force-local-datapipe", action="store_true") |
| args = parser.parse_args() |
|
|
| dist = initialize_distributed() |
| device = dist.device |
| cfg = prepare_config(load_config(args.config), data_dir=args.data_dir, checkpoint=args.checkpoint) |
| if args.output_dir: |
| cfg.inference.output_dir = str(Path(args.output_dir).expanduser().resolve()) |
|
|
| checkpoint = Path(args.checkpoint or cfg.inference.checkpoint) |
| if not checkpoint.is_file(): |
| raise FileNotFoundError(f"checkpoint not found: {checkpoint}") |
|
|
| datapipe = build_datapipe( |
| cfg, |
| distributed=False, |
| force_local=args.force_local_datapipe, |
| ) |
| val_loader, _ = datapipe.val_dataloader() |
| model = build_model(datapipe.spatial_dim, cfg).to(device) |
| model.load_state_dict(load_model_state(checkpoint, device)) |
| model.eval() |
|
|
| output_dir = Path(args.output_dir or cfg.inference.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| max_batches = args.max_batches or int(get_attr(cfg.inference, "max_batches", 1)) |
|
|
| written = 0 |
| for batch_idx, (x, y, grid) in enumerate(val_loader): |
| if batch_idx >= max_batches: |
| break |
| x = x.to(device) |
| y = y.to(device) |
| grid = grid.to(device) |
| pred, target = predict_batch(model, x, y, grid, cfg) |
| path = output_dir / f"batch_{batch_idx:04d}.npz" |
| np.savez_compressed( |
| path, |
| prediction=pred.detach().cpu().numpy(), |
| target=target.detach().cpu().numpy(), |
| ) |
| print(f"wrote {path}") |
| written += 1 |
|
|
| if written == 0: |
| raise RuntimeError("no validation batches were available for inference") |
|
|
| cleanup_distributed() |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
