File size: 3,010 Bytes
5978445 | 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 | """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()
|