| """Load the saved best FNO checkpoint and evaluate untouched official test data.""" |
| from pathlib import Path |
|
|
| from common import (case_guard, config_hash, load_config, load_data, |
| negative_total_flow_rates, parse_args, predict, save_json, |
| save_torch, sha256) |
| import torch |
| from model.NOTS import NOTS |
|
|
|
|
| def main(): |
| parser = parse_args(__doc__) |
| parser.add_argument('--checkpoint', type=Path) |
| args = parser.parse_args() |
| cfg = load_config(args.config) |
| device = case_guard(cfg) |
| out = Path(cfg['paths']['output_dir']) |
| checkpoint_path = args.checkpoint or out / 'weight/best_model.pt' |
| state = torch.load(checkpoint_path, weights_only=False, map_location='cpu') |
| if state['config_hash'] != config_hash(cfg): |
| raise ValueError('Checkpoint/config protocol mismatch') |
| if state['normalization']['kind'] != 'identity': |
| raise ValueError('Unexpected checkpoint prediction scale') |
| model = NOTS(**state['model_config']).to(device) |
| model.load_state_dict(state['model_state'], strict=True) |
| model.eval() |
| metrics = {'checkpoint': str(checkpoint_path), 'checkpoint_sha256': sha256(checkpoint_path), |
| 'selected_seed': state['seed'], 'selection_train_mse': state['train_mse'], 'splits': {}} |
| for split in ('test', 'train'): |
| x, y, meta = load_data(cfg, split) |
| if split == 'train' and meta['sha256'] != state['data_sha256']: |
| raise ValueError('Checkpoint training pool hash differs') |
| x, y = x.to(device), y.to(device) |
| with torch.no_grad(): |
| output = predict(model, x, cfg['experiment']['prediction_batch_size']) |
| error = output - y |
| truth = negative_total_flow_rates(y, x) |
| predicted_flow = negative_total_flow_rates(output, x) |
| result = {'n': len(x), 'mse': error.square().mean().item(), |
| 'relative_l2_mean': (error.flatten(1).norm(dim=1) / y.flatten(1).norm(dim=1).clamp_min(1e-12)).mean().item(), |
| 'functional_mae': (truth - predicted_flow).abs().mean().item(), |
| 'predicted_best_index': predicted_flow.argmax().item(), |
| 'true_best_index': truth.argmax().item(), |
| 'prediction_shape': list(output.shape), |
| 'all_finite': bool(torch.isfinite(output).all().item()), 'data': meta} |
| if not result['all_finite']: |
| raise FloatingPointError('Nonfinite checkpoint inference') |
| metrics['splits'][split] = result |
| save_torch(out / 'results' / f'inference_{split}.pt', {'x': x.cpu(), 'target': y.cpu(), |
| 'prediction': output.cpu(), 'true_functional': truth.cpu(), |
| 'predicted_functional': predicted_flow.cpu(), 'checkpoint_sha256': metrics['checkpoint_sha256']}) |
| print(f'INFERENCE split={split} {result}', flush=True) |
| save_json(out / 'results/inference_metrics.json', metrics) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
