File size: 4,308 Bytes
8f89e30 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Reload the best checkpoint and evaluate every held-out validation sample."""
import argparse
import csv
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from model.lino import LiNO
from common import (compute_device, decode, file_digest, load_checkpoint, load_config,
make_inputs, package_path, read_fields, relative_l2, save_json)
@torch.no_grad()
def inference(config, checkpoint=None):
device = compute_device(config["inference"]["device"])
checkpoint = package_path(checkpoint or config["training"]["checkpoint"])
state = load_checkpoint(checkpoint)
contract = state["data_contract"]
for key in ["resolution", "interpolation", "input_key", "target_key"]:
if config["data"][key] != contract[key]:
raise ValueError(f"Data contract mismatch for {key}")
data_path = package_path(config["data"]["directory"]) / config["data"]["filename"]
if file_digest(data_path) != contract["dataset_sha256"]:
raise ValueError("Dataset content differs from the training data (paths may change)")
initial, final = read_fields(config)
ids, stats = state["val_indices"], state["normalizers"]
model = LiNO(**state["architecture"]).to(device)
model.load_state_dict(state["model_state"], strict=True)
model.eval()
inputs = make_inputs(initial[ids], stats)
targets = final[ids]
predictions = []
started = time.monotonic()
batch_size = config["inference"]["batch_size"]
for start in range(0, len(ids), batch_size):
normalized = model(inputs[start:start+batch_size].to(device))
predictions.append(decode(normalized, stats).cpu())
if start % (batch_size * 25) == 0:
print(f"INFERENCE samples={min(start+batch_size,len(ids))}/{len(ids)} checkpoint_epoch={state['epoch']}", flush=True)
prediction = torch.cat(predictions)
errors = relative_l2(prediction, targets)
mse_by_sample = (prediction-targets).square().flatten(1).mean(1)
if not torch.isfinite(prediction).all():
raise FloatingPointError("Nonfinite inference output")
output = package_path(config["results"]["directory"])
output.mkdir(parents=True, exist_ok=True)
np.savez_compressed(output/"predictions.npz", initial=initial[ids].numpy(), target=targets.numpy(),
prediction=prediction.numpy(), sample_indices=ids.numpy(), relative_l2=errors.numpy(),
x=np.linspace(0,1,config["data"]["resolution"],dtype=np.float32))
with open(output/"per_sample_metrics.csv", "w", newline="") as stream:
writer = csv.writer(stream)
writer.writerow(["validation_index", "source_index", "relative_l2", "mse_physical"])
writer.writerows((i,int(source),float(rel),float(mse)) for i,(source,rel,mse) in enumerate(zip(ids,errors,mse_by_sample)))
metric = errors.mean().item()
metrics = dict(n_validation=len(ids), checkpoint_epoch=state["epoch"], best_epoch=state["best_epoch"],
validation_relative_l2=metric, validation_relative_l2_std=errors.std(correction=0).item(),
validation_mse_physical=mse_by_sample.mean().item(),
validation_rmse_physical=mse_by_sample.mean().sqrt().item(),
checkpoint_recorded_relative_l2=state["best_val_relative_l2"],
reload_metric_absolute_difference=abs(metric-state["best_val_relative_l2"]),
paper_table1_relative_l2=0.0109, ratio_to_paper=metric/0.0109,
inference_seconds=time.monotonic()-started,
checkpoint_sha256=file_digest(checkpoint), device=torch.cuda.get_device_name(device))
if metrics["reload_metric_absolute_difference"] > 1e-6:
raise AssertionError(f"Reloaded metric disagrees with checkpoint: {metrics}")
save_json(output/"metrics.json", metrics)
print("BEST CHECKPOINT INFERENCE COMPLETED", metrics, flush=True)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", default="config/config.yaml")
parser.add_argument("--checkpoint", default=None)
args = parser.parse_args()
inference(load_config(args.config), args.checkpoint)
|