File size: 3,790 Bytes
0d2f59e | 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 | """Load the selected checkpoint and query the three paper snapshot grids."""
from pathlib import Path
import sys
import time
import numpy as np
import torch
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from model import PINNShallowWater
from common import arguments, device_for, errors, load_config, output_path, runtime_info, save_json, sha256
from fake_data import initial_values
from train import to_device, validation, weights
def main():
args = arguments(__doc__).parse_args()
config, root = load_config(args.config)
device = device_for(config)
torch.set_num_threads(8)
checkpoint_path = output_path(root, config['paths']['best_checkpoint'])
checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
if checkpoint.get('completed_iterations') != 5000 or not checkpoint.get('formal_run_complete'):
raise ValueError('Formal 5000-iteration training has not completed.')
if checkpoint['config'] != config or checkpoint['model_class'] != 'PINNShallowWater':
raise ValueError('Checkpoint/config/model identity mismatch.')
model = PINNShallowWater(**config['model']).to(device)
model.load_state_dict(checkpoint['model_state_dict'], strict=True)
model.eval()
with np.load(output_path(root, config['paths']['grid'])) as grid:
x, y, xy, times = [grid[k] for k in ['x', 'y', 'xy', 'times']]
batch_size = config['evaluation']['inference_batch_size']
snapshots = []
started = time.monotonic()
with torch.no_grad():
for t in times:
coordinates = np.column_stack((xy, np.full(len(xy), t))).astype(np.float32)
batches = []
for begin in range(0, len(coordinates), batch_size):
prediction = model(torch.as_tensor(coordinates[begin:begin + batch_size], device=device))
batches.append(prediction.cpu().numpy())
field = np.concatenate(batches).reshape(len(y), len(x), 3)
if not np.isfinite(field).all():
raise FloatingPointError(f'Nonfinite prediction at t={t}')
snapshots.append(field)
print('INFERENCE_SNAPSHOT', {'t': float(t), 'shape': field.shape,
'min_huv': field.reshape(-1, 3).min(axis=0).tolist(),
'max_huv': field.reshape(-1, 3).max(axis=0).tolist()}, flush=True)
fields = np.stack(snapshots)
path = output_path(root, config['paths']['predictions'])
np.savez_compressed(path, x=x, y=y, times=times, fields=fields,
variables=np.array(['h', 'u', 'v']), units=np.array(['m', 'm/s', 'm/s']))
initial_errors = errors(fields[0].reshape(-1, 3), initial_values(xy, config['physics']))
samples_path = output_path(root, config['paths']['samples'])
if checkpoint['sampling_sha256'] != sha256(samples_path):
raise ValueError('Checkpoint sampling hash mismatch.')
with np.load(samples_path) as arrays:
val = to_device(arrays, 'val_', device)
check = validation(model, val, config, weights(config, device))
metadata = {'status': 'completed', 'checkpoint': str(checkpoint_path),
'checkpoint_sha256': sha256(checkpoint_path), 'selected_iteration': checkpoint['iteration'],
'completed_training_iterations': 5000, 'state_dict_strict_load': True,
'fields_shape': list(fields.shape), 'all_finite': True, 'initial_grid_errors': initial_errors,
'best_checkpoint_validation': check, 'runtime': runtime_info(device),
'elapsed_seconds': time.monotonic() - started, 'predictions_sha256': sha256(path)}
save_json(path.parent / 'inference_summary.json', metadata)
print('INFERENCE_COMPLETE', metadata, flush=True)
if __name__ == '__main__':
main()
|