File size: 5,551 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Data, normalization, metrics and checkpoint contracts shared by the CLIs."""
from __future__ import annotations
import argparse
import hashlib
import json
import os
from pathlib import Path
import random
import sys
import numpy as np
from scipy.io import loadmat
import torch
import torch.nn.functional as F
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.hno import HNO
def configuration(description):
parser = argparse.ArgumentParser(description=description)
parser.add_argument('--config', default='config/config.yaml')
parser.add_argument('--device', default='cuda', choices=['cuda', 'cpu'])
return parser
def read_config(path):
with open(path, encoding='utf-8') as f:
c = yaml.safe_load(f)
for section in ['model', 'data', 'training', 'output']:
if section not in c:
raise ValueError(f'Missing config section: {section}')
if c['model']['grid_size'] != (c['data']['raw_shape'][1] - 1) // c['data']['stride'] + 1:
raise ValueError('Configured grid and strided raw data disagree')
if c['training']['precision'] != 'float32':
raise ValueError('This reproduction uses FP32 model parameters')
return c
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def compute_device(name):
if name == 'cuda':
if not os.environ.get('SLURM_JOB_ID'):
raise RuntimeError('Accelerator computation requires a Slurm allocation')
if not torch.cuda.is_available():
raise RuntimeError('No accelerator visible in this Slurm allocation')
return torch.device(name)
def read_split(c, split):
d = c['data']
count = d['n_train' if split == 'train' else 'n_eval']
path = Path(d['root']) / d['train_file' if split == 'train' else 'eval_file']
print(f'Reading {split}: {path}, samples={count}', flush=True)
raw = loadmat(path, variable_names=['coeff', 'sol'])
arrays = []
for key in ['coeff', 'sol']:
a = raw.pop(key)
if list(a.shape) != d['raw_shape']:
raise ValueError(f'{path}:{key} shape {a.shape}, expected {d["raw_shape"]}')
if not 0 < count <= len(a):
raise ValueError('Invalid sample count')
selected = np.array(a[:count, ::d['stride'], ::d['stride']], copy=True, order='C')
if not np.isfinite(selected).all():
raise ValueError(f'Nonfinite {key} data')
tensor = torch.from_numpy(selected)
if key == 'coeff':
tensor = tensor.float()
arrays.append(tensor)
del a
if arrays[0].shape[1:] != (c['model']['grid_size'], c['model']['grid_size']):
raise ValueError('Unexpected downsampled grid')
return tuple(arrays)
def training_statistics(coeff, solution):
# Keep non-scalar tensor statistics: preserves source solution precision.
return {name: value.reshape(1, 1, 1) for name, value in {
'input_mean': coeff.mean(), 'input_std': coeff.std() + 1e-8,
'target_mean': solution.mean(), 'target_std': solution.std() + 1e-8,
}.items()}
def encode_input(x, stats):
return ((x - stats['input_mean']) / stats['input_std']).unsqueeze(-1)
def decode_output(y, stats):
return y[..., 0] * stats['target_std'] + stats['target_mean']
def relative_l2(prediction, target):
denominator = torch.linalg.vector_norm(target.flatten(1), dim=1)
return torch.linalg.vector_norm((prediction - target).flatten(1), dim=1) / denominator.clamp_min(1e-12)
def central_gradient(field, spacing):
extended = F.pad(field, (1, 1, 1, 1))
dx = (extended[:, 1:-1, 2:] - extended[:, 1:-1, :-2]) / (2 * spacing)
dy = (extended[:, 2:, 1:-1] - extended[:, :-2, 1:-1]) / (2 * spacing)
return dx, dy
def objective(prediction, target, training):
primary = relative_l2(prediction, target).sum()
interior_prediction = F.pad(prediction[:, 1:-1, 1:-1], (1, 1, 1, 1))
predicted_gradients = central_gradient(interior_prediction, training['derivative_spacing'])
target_gradients = central_gradient(target, training['derivative_spacing'])
auxiliary = sum(relative_l2(p, t).sum() for p, t in zip(predicted_gradients, target_gradients))
return primary + training['derivative_weight'] * auxiliary, primary, auxiliary
def atomic_json(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix(path.suffix + '.tmp')
temporary.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False), encoding='utf-8')
temporary.replace(path)
def atomic_checkpoint(path, value):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_suffix('.tmp')
torch.save(value, temporary)
temporary.replace(path)
def load_checkpoint(path, device):
checkpoint = torch.load(path, map_location='cpu', weights_only=True)
if checkpoint['format_version'] != 1:
raise ValueError('Unsupported checkpoint schema')
model = HNO(**checkpoint['model_config']).to(device)
model.load_state_dict(checkpoint['model_state'], strict=True)
stats = {k: v.to(device) for k, v in checkpoint['normalization'].items()}
return model, stats, checkpoint
def source_hashes():
paths = [ROOT/'model/hno.py', *sorted((ROOT/'scripts').glob('*.py'))]
return {str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in paths}
|