| """Full 5000-update dCDB9 experiment; no reduced-training mode.""" |
| import csv |
| import os |
| from pathlib import Path |
| import random |
| 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, checkpoint_save, device_for, errors, load_config, |
| output_path, runtime_info, save_json, sha256) |
| from fake_data import generate |
|
|
|
|
| def check_formal_config(config): |
| required = [config['experiment'] == 'dCDB9', |
| config['model']['hidden_layers'] == 5, config['model']['width'] == 300, |
| config['training']['iterations'] == 5000, |
| config['training']['learning_rate'] == 0.001, |
| config['training']['scheduler'] is None, |
| config['training']['early_stopping'] is False, |
| [config['sampling'][k] for k in ('interior', 'boundary', 'initial')] == [32768, 8192, 8192], |
| config['loss']['pde'] == [10, 1, 1], |
| config['loss']['initial'] == [100, 100, 100], |
| config['loss']['positive_depth'] == config['loss']['dry_velocity'] == 10, |
| config['loss']['entropy'] == 0, |
| config['physics']['t_range'] == [0, 0.8]] |
| if not all(required): |
| raise ValueError('Configuration violates the formal paper/user dCDB9 experiment.') |
|
|
|
|
| def residual_terms(model, x, gravity, create_graph): |
| state, residual = model.residuals(x, gravity, create_graph=create_graph) |
| mask = (state[:, 0] < 0).to(state.dtype) |
| return torch.cat((residual.square().mean(dim=0), |
| (state.square() * mask[:, None]).mean(dim=0))) |
|
|
|
|
| def boundary_initial_terms(model, data): |
| ni, nb = len(data['initial']), len(data['boundary']) |
| |
| values = model(torch.cat((data['initial'], data['boundary'], data['boundary_pair']))) |
| initial = (values[:ni] - data['initial_target']).square().mean(dim=0) |
| diff = values[ni:ni + nb] - values[ni + nb:] |
| |
| boundary = diff.reshape(4, nb // 4, 3).square().mean(dim=1).sum(dim=0) |
| return initial, boundary |
|
|
|
|
| def weights(config, device): |
| loss = config['loss'] |
| return (torch.tensor(loss['pde'] + [loss['positive_depth'], loss['dry_velocity'], loss['dry_velocity']], device=device), |
| torch.tensor(loss['initial'], device=device), |
| torch.tensor(loss['boundary'], device=device)) |
|
|
|
|
| def validation(model, data, config, weight_vectors): |
| wp, wi, wb = weight_vectors |
| terms = torch.zeros(6, device=wp.device) |
| n = len(data['interior']) |
| chunk = min(4096, config['training']['pde_chunk_size']) |
| with torch.enable_grad(): |
| for start in range(0, n, chunk): |
| batch = data['interior'][start:start + chunk].detach().requires_grad_(True) |
| terms += residual_terms(model, batch, config['physics']['gravity'], False).detach() * (len(batch) / n) |
| with torch.no_grad(): |
| li, lb = boundary_initial_terms(model, data) |
| recovery = model.slope_recovery() |
| value = terms @ wp + li @ wi + lb @ wb + config['loss']['slope_recovery'] * recovery |
| target = data['initial_target'] |
| diff = model(data['initial']) - target |
| ic_mae = diff.abs().mean(dim=0).cpu().tolist() |
| ic_rmse = diff.square().mean(dim=0).sqrt().cpu().tolist() |
| return {'objective': value.item(), 'pde_mse': terms[:3].cpu().tolist(), |
| 'constraint_mse': terms[3:].cpu().tolist(), 'boundary_mse_sum_per_variable': lb.cpu().tolist(), |
| 'ic_mae': ic_mae, 'ic_rmse': ic_rmse, 'slope_recovery': recovery.item()} |
|
|
|
|
| def to_device(arrays, prefix, device): |
| return {key: torch.as_tensor(arrays[prefix + key], device=device) |
| for key in ['interior', 'initial', 'initial_target', 'boundary', 'boundary_pair']} |
|
|
|
|
| def main(): |
| parser = arguments(__doc__) |
| parser.add_argument('--resume', type=Path) |
| parser.add_argument('--check-config', action='store_true') |
| args = parser.parse_args() |
| config, root = load_config(args.config) |
| check_formal_config(config) |
| if args.check_config: |
| print('FORMAL_CONFIG_OK: 5x300, 32768/8192/8192 points, 5000 Adam updates, no decay/early stopping.') |
| return |
| device = device_for(config) |
| torch.set_num_threads(min(int(os.environ.get('SLURM_CPUS_PER_TASK', '8')), 8)) |
| seed = config['seed'] |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| samples_path = output_path(root, config['paths']['samples']) |
| if not samples_path.exists(): |
| generate(config, root) |
| sample_hash = sha256(samples_path) |
| with np.load(samples_path) as arrays: |
| train = to_device(arrays, '', device) |
| val = to_device(arrays, 'val_', device) |
| for key, expected in [('interior', 32768), ('initial', 8192), ('boundary', 8192)]: |
| assert train[key].shape == (expected, 3), (key, train[key].shape) |
| model = PINNShallowWater(**config['model']).to(device) |
| optimizer = torch.optim.Adam(model.parameters(), lr=config['training']['learning_rate'], |
| betas=tuple(config['training']['betas']), eps=config['training']['epsilon'], |
| weight_decay=config['training']['weight_decay']) |
| wp, wi, wb = weight_vectors = weights(config, device) |
| best_path = output_path(root, config['paths']['best_checkpoint']) |
| last_path = output_path(root, config['paths']['last_checkpoint']) |
| result_dir = output_path(root, config['paths']['results']) |
| result_dir.mkdir(parents=True, exist_ok=True) |
| history_path = result_dir / 'training_history.csv' |
| best, best_step, first, prior_seconds = float('inf'), 0, 1, 0.0 |
| hardware = runtime_info(device) |
| if args.resume: |
| checkpoint = torch.load(args.resume, map_location=device, weights_only=False) |
| if checkpoint['config'] != config or checkpoint['sampling_sha256'] != sample_hash: |
| raise ValueError('Resume config or sampling data differs from the saved experiment.') |
| model.load_state_dict(checkpoint['model_state_dict'], strict=True) |
| optimizer.load_state_dict(checkpoint['optimizer_state_dict']) |
| first = checkpoint['iteration'] + 1 |
| best, best_step = checkpoint['best_validation_objective'], checkpoint['best_iteration'] |
| prior_seconds = checkpoint['elapsed_seconds'] |
| torch.set_rng_state(checkpoint['torch_rng'].cpu()) |
| torch.cuda.set_rng_state_all([x.cpu() for x in checkpoint['cuda_rng']]) |
| np.random.set_state(checkpoint['numpy_rng']) |
| random.setstate(checkpoint['python_rng']) |
| if history_path.exists(): |
| with history_path.open() as source: |
| rows = list(csv.DictReader(source)) |
| rows = [r for r in rows if int(r['iteration']) < first] |
| with history_path.open('w', newline='') as dest: |
| writer = csv.DictWriter(dest, fieldnames=list(rows[0]) if rows else fieldnames()) |
| writer.writeheader() |
| writer.writerows(rows) |
| elif best_path.exists() or last_path.exists(): |
| raise FileExistsError('Existing experiment weights found. Use --resume explicitly.') |
| print('FORMAL_TRAINING_START', {'runtime': hardware, 'parameters': sum(p.numel() for p in model.parameters()), |
| 'sampling_sha256': sample_hash, 'iterations': 5000, 'starting_iteration': first, |
| 'samples_every_update': [32768, 8192, 8192], 'external_training_data': False}, flush=True) |
| save_json(result_dir / 'runtime_environment.json', hardware) |
| start_time = time.monotonic() |
| fields = fieldnames() |
| mode = 'a' if args.resume and history_path.exists() else 'w' |
| with history_path.open(mode, newline='', buffering=1) as handle: |
| writer = csv.DictWriter(handle, fieldnames=fields) |
| if mode == 'w': |
| writer.writeheader() |
| last_validation = None |
| for iteration in range(first, config['training']['iterations'] + 1): |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| pde_terms = torch.zeros(6, device=device) |
| n = len(train['interior']) |
| chunk = config['training']['pde_chunk_size'] |
| for begin in range(0, n, chunk): |
| points = train['interior'][begin:begin + chunk].detach().requires_grad_(True) |
| part = residual_terms(model, points, config['physics']['gravity'], True) |
| fraction = len(points) / n |
| (fraction * (part @ wp)).backward() |
| pde_terms += fraction * part.detach() |
| initial_loss, boundary_loss = boundary_initial_terms(model, train) |
| recovery = model.slope_recovery() |
| remaining = initial_loss @ wi + boundary_loss @ wb + config['loss']['slope_recovery'] * recovery |
| remaining.backward() |
| loss = pde_terms @ wp + remaining.detach() |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f'Nonfinite loss at iteration {iteration}') |
| optimizer.step() |
| values = torch.cat((loss.reshape(1), pde_terms, initial_loss.detach(), |
| boundary_loss.detach(), recovery.detach().reshape(1))).cpu().tolist() |
| elapsed = time.monotonic() - start_time + prior_seconds |
| row = dict(zip(fields[:15], [iteration] + values)) |
| row.update(learning_rate=optimizer.param_groups[0]['lr'], elapsed_seconds=elapsed) |
| writer.writerow(row) |
| if iteration % config['training']['print_every'] == 0: |
| print('TRAIN', row, flush=True) |
| if iteration == 1 or iteration % config['training']['validation_every'] == 0 or iteration == 5000: |
| model.eval() |
| last_validation = validation(model, val, config, weight_vectors) |
| print('EVAL', {'iteration': iteration, **last_validation}, flush=True) |
| save_json(result_dir / f'validation_latest.json', {'iteration': iteration, **last_validation}) |
| if last_validation['objective'] < best: |
| best, best_step = last_validation['objective'], iteration |
| checkpoint_save(best_path, {'model_class': 'PINNShallowWater', 'model_state_dict': |
| {k: v.detach().cpu() for k, v in model.state_dict().items()}, |
| 'iteration': iteration, 'completed_iterations': iteration, 'config': config, |
| 'validation': last_validation, 'sampling_sha256': sample_hash, 'runtime': hardware}) |
| if iteration % config['training']['checkpoint_every'] == 0 or iteration == 5000: |
| checkpoint_save(last_path, {'model_class': 'PINNShallowWater', 'model_state_dict': model.state_dict(), |
| 'optimizer_state_dict': optimizer.state_dict(), 'iteration': iteration, 'completed_iterations': iteration, |
| 'best_iteration': best_step, 'best_validation_objective': best, 'config': config, |
| 'sampling_sha256': sample_hash, 'runtime': hardware, 'elapsed_seconds': elapsed, |
| 'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(), |
| 'numpy_rng': np.random.get_state(), 'python_rng': random.getstate()}) |
| save_json(result_dir / 'training_progress.json', {'completed_iterations': iteration, |
| 'required_iterations': 5000, 'best_iteration': best_step, |
| 'best_validation_objective': best, 'elapsed_seconds': elapsed}) |
| handle.flush() |
| completed = torch.load(last_path, map_location='cpu', weights_only=False) |
| assert completed['iteration'] == 5000 |
| best_checkpoint = torch.load(best_path, map_location='cpu', weights_only=False) |
| best_checkpoint['completed_iterations'] = 5000 |
| best_checkpoint['formal_run_complete'] = True |
| checkpoint_save(best_path, best_checkpoint) |
| summary = {'status': 'completed', 'completed_iterations': 5000, 'optimizer_updates': 5000, |
| 'stopped_early': False, 'best_iteration': best_step, 'best_validation_objective': best, |
| 'last_validation': last_validation, 'elapsed_seconds': elapsed, 'runtime': hardware, |
| 'best_checkpoint_sha256': sha256(best_path), 'last_checkpoint_sha256': sha256(last_path), |
| 'sampling_sha256': sample_hash, 'history_rows': 5000, |
| 'external_training_data': False, 'samples_every_update': [32768, 8192, 8192]} |
| save_json(result_dir / 'training_summary.json', summary) |
| print('FORMAL_TRAINING_COMPLETE', summary, flush=True) |
|
|
|
|
| def fieldnames(): |
| return ['iteration', 'loss', 'mass_mse', 'momentum_x_mse', 'momentum_y_mse', |
| 'positive_depth_mse', 'dry_u_mse', 'dry_v_mse', 'initial_h_mse', 'initial_u_mse', |
| 'initial_v_mse', 'boundary_h_mse', 'boundary_u_mse', 'boundary_v_mse', |
| 'slope_recovery', 'learning_rate', 'elapsed_seconds'] |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
