File size: 1,302 Bytes
986404c | 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 | import argparse
from pathlib import Path
import numpy as np
import torch
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model.corrdiff import CorrDiff
def main():
parser = argparse.ArgumentParser(description="Run CorrDiff inference")
parser.add_argument("--data", default="data/era5_corrdiff.npz")
parser.add_argument("--checkpoint", default="data/checkpoints/model_bak.pth")
parser.add_argument("--ensemble-size", type=int, default=1)
parser.add_argument("--output", default="result/output/predictions.npz")
args = parser.parse_args()
data = np.load(args.data)
coarse = torch.from_numpy(data["input"])
model = CorrDiff()
checkpoint = Path(args.checkpoint)
if checkpoint.exists():
model.load_state_dict(torch.load(checkpoint, map_location="cpu")["model"])
model.eval()
samples = []
with torch.no_grad():
for _ in range(args.ensemble_size):
samples.append(model(coarse).numpy().astype("float32"))
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output, prediction=np.stack(samples), input=data["input"])
print(f"prediction: {np.stack(samples).shape}")
print(f"saved: {output}")
if __name__ == "__main__":
main()
|