File size: 4,257 Bytes
09c43eb | 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 | """Evaluate all 200 Darcy examples with the independently loaded best weights."""
from __future__ import annotations
import csv
from pathlib import Path
import time
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset
from common import (atomic_json, compute_device, configuration, decode_output,
encode_input, load_checkpoint, read_config, read_split, relative_l2)
def main():
parser=configuration('HNO Darcy best-checkpoint inference')
parser.add_argument('--checkpoint', help='Override the checkpoint path in config')
args=parser.parse_args()
c=read_config(args.config)
device=compute_device(args.device)
torch.set_num_threads(c['training']['cpu_threads'])
torch.backends.cudnn.benchmark=True
checkpoint_path=Path(args.checkpoint or c['output']['checkpoint'])
model,stats,checkpoint=load_checkpoint(checkpoint_path,device)
if checkpoint['model_config']!=c['model']:
raise ValueError('Current model config disagrees with saved architecture')
model.eval()
print(f'Loaded best weights: {checkpoint_path}; best epoch={checkpoint["epoch"]}; '
f'selection relative L2={checkpoint["eval_relative_l2"]:.9f}',flush=True)
coefficient,target=read_split(c,'eval')
cpu_stats={k:v.cpu() for k,v in stats.items()}
loader=DataLoader(TensorDataset(encode_input(coefficient,cpu_stats),target),
batch_size=c['inference']['batch_size'],shuffle=False)
predicted=[]; relative=[]; elapsed=0
with torch.no_grad():
for index,(x,y) in enumerate(loader,start=1):
x,y=x.to(device),y.to(device)
if device.type=='cuda': torch.cuda.synchronize()
tic=time.perf_counter()
p=decode_output(model(x),stats)
if device.type=='cuda': torch.cuda.synchronize()
elapsed+=time.perf_counter()-tic
r=relative_l2(p,y)
predicted.append(p.cpu().numpy()); relative.extend(r.cpu().tolist())
print(f'inference batch={index:02d}/{len(loader)} mean_relative_l2={r.mean().item():.8f}',flush=True)
prediction=np.concatenate(predicted)
truth=target.numpy()
if prediction.shape!=truth.shape or not np.isfinite(prediction).all():
raise ValueError('Invalid prediction output')
error=prediction-truth
output=Path(c['output']['results_dir']);output.mkdir(parents=True,exist_ok=True)
np.savez_compressed(output/'predictions.npz',coefficient=coefficient.numpy(),target=truth,
prediction=prediction,relative_l2=np.asarray(relative),
sample_indices=np.arange(len(truth)),checkpoint_epoch=checkpoint['epoch'])
with (output/'per_sample_metrics.csv').open('w',newline='') as f:
writer=csv.writer(f);writer.writerow(['sample_index','relative_l2','rmse','mae'])
for i,r in enumerate(relative):
writer.writerow([i,r,float(np.sqrt(np.mean(error[i]**2))),float(np.mean(np.abs(error[i])))])
metrics={'n_eval':len(relative),'resolution':list(truth.shape[1:]),
'mean_relative_l2':float(np.mean(relative)), 'std_relative_l2':float(np.std(relative)),
'median_relative_l2':float(np.median(relative)), 'min_relative_l2':float(np.min(relative)),
'max_relative_l2':float(np.max(relative)), 'rmse':float(np.sqrt(np.mean(error**2))),
'mae':float(np.mean(np.abs(error))), 'best_epoch':checkpoint['epoch'],
'checkpoint_eval_relative_l2':checkpoint['eval_relative_l2'],
'checkpoint':str(checkpoint_path),'inference_seconds':elapsed,
'milliseconds_per_batch':elapsed/len(loader)*1000,
'paper_relative_l2':c['reference']['paper_relative_l2'],
'relative_l2_ratio_to_paper':float(np.mean(relative))/c['reference']['paper_relative_l2'],
'protocol_note':c['reference']['note']}
if not np.isclose(metrics['mean_relative_l2'],checkpoint['eval_relative_l2'],rtol=5e-4,atol=1e-7):
raise ValueError('Loaded best-checkpoint metric does not reproduce saved evaluation')
atomic_json(output/'metrics.json',metrics)
print('INFERENCE_COMPLETED '+str(metrics),flush=True)
if __name__=='__main__':
main()
|