| """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() |
|
|
