| |
| """Paper Table 1, Re=800: 100000-step adaptive pseudo-time training.""" |
| import argparse |
| import copy |
| import csv |
| import hashlib |
| import json |
| import os |
| from pathlib import Path |
| import random |
| import socket |
| import sys |
| import time |
| import traceback |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import torch |
|
|
| from model.soap import SOAP |
| from scripts.common import (LOSS_NAMES, load_config, atomic_json, atomic_save, |
| sample_interior, analytic_boundaries, boundary_losses, |
| read_reference, predict, errors, build_model, |
| require_compute_node) |
|
|
|
|
| def learning_rate(index, cfg): |
| o = cfg['optimizer'] |
| warm = o['warmup_steps'] |
| if index < warm: |
| return o['learning_rate'] * index / warm |
| return o['learning_rate'] * o['decay_rate'] ** ((index - warm) / o['decay_steps']) |
|
|
|
|
| @torch.no_grad() |
| def update_tau(tau, value, previous_value, residual, previous_residual, initial_loss, cfg): |
| p = cfg['pseudo_time'] |
| eps = p['epsilon'] |
| loss = residual.square().mean(0) |
| reduction = torch.log10((initial_loss + eps) / (loss + eps)) |
| progress = ((reduction - p['shrink_start']) / |
| (p['shrink_end'] - p['shrink_start'])).clamp(0.0, 1.0) |
| gamma = p['minimum_shrink'] + (1.0 - p['minimum_shrink']) * (1.0 + torch.cos(torch.pi * progress)) / 2.0 |
| tau_hat = gamma * torch.linalg.vector_norm(value - previous_value, dim=0) / ( |
| torch.linalg.vector_norm(residual - previous_residual, dim=0) + eps) |
| |
| return ((1.0 - p['beta']) * tau + p['beta'] * tau_hat).detach(), gamma |
|
|
|
|
| def gradnorm_weights(losses, parameters, weights, cfg): |
| norms = [] |
| for loss in losses.unbind(): |
| grads = torch.autograd.grad(loss, parameters, retain_graph=True, allow_unused=True) |
| total = sum(g.square().sum() for g in grads if g is not None) |
| norms.append(total.sqrt()) |
| norms = torch.stack(norms).detach() |
| w = cfg['loss_weighting'] |
| |
| proposed = norms.sum() / norms.clamp_min(w['epsilon']) |
| return (w['momentum'] * weights + (1.0 - w['momentum']) * proposed).detach(), norms |
|
|
|
|
| @torch.no_grad() |
| def physical_validation(model, xy, boundary, viscosity): |
| _, r = model.physics(xy, viscosity) |
| bl = boundary_losses(model, boundary) |
| return float(r.square().mean(0).sum() + bl.sum()) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--config', default='config/config.yaml') |
| parser.add_argument('--resume', type=Path) |
| args = parser.parse_args() |
| cfg, root = load_config(args.config) |
| results = root / cfg['paths']['results'] |
| results.mkdir(parents=True, exist_ok=True) |
| require_compute_node(cfg) |
| device = torch.device(cfg['training']['device']) |
| torch.set_num_threads(min(int(os.environ.get('SLURM_CPUS_PER_TASK', '8')), 8)) |
| random.seed(cfg['seed']) |
| np.random.seed(cfg['seed']) |
| torch.manual_seed(cfg['seed']) |
| torch.cuda.manual_seed_all(cfg['seed']) |
| model = build_model(cfg, device) |
| previous = copy.deepcopy(model).requires_grad_(False) |
| optcfg = cfg['optimizer'] |
| optimizer = SOAP(model.parameters(), lr=0.0, |
| betas=(optcfg['beta1'], optcfg['beta2']), eps=optcfg['epsilon'], |
| shampoo_beta=optcfg['shampoo_beta'], |
| precondition_frequency=optcfg['precondition_frequency']) |
| if cfg['training']['jit']: |
| model = torch.jit.script(model) |
| previous = torch.jit.script(previous) |
| parameters = tuple(p for p in model.parameters() if p.requires_grad) |
| nu = cfg['domain']['viscosity'] |
| boundary = analytic_boundaries(cfg, device) |
| |
| reference_xy, reference_uvp, data_meta, _ = read_reference( |
| cfg['evaluation']['reference_vtu'], cfg['evaluation']['expected_points']) |
| atomic_json(results / 'reference_metadata.json', data_meta) |
| validation_generator = torch.Generator(device=device).manual_seed(cfg['training']['validation_seed']) |
| validation_xy = sample_interior(cfg, device, validation_generator, cfg['training']['validation_points']) |
| tau = torch.full((3,), cfg['pseudo_time']['initial_tau'], device=device) |
| weights = torch.tensor([cfg['loss_weighting']['initial_boundary']] * 5 + |
| [cfg['loss_weighting']['initial_residual']] * 3, device=device) |
| initial_loss = None |
| gamma = torch.ones(3, device=device) |
| best_score, best_step, completed = float('inf'), 0, 0 |
| requested = cfg['training']['steps'] |
| if requested != 100000: |
| raise ValueError('Formal reproduction requires exactly 100000 steps') |
| start = time.monotonic() |
| previous_seconds = 0.0 |
| if args.resume: |
| ck = torch.load(args.resume, map_location=device, weights_only=False) |
| if ck['config']['model'] != cfg['model']: |
| raise ValueError('Resume model configuration differs') |
| model.load_state_dict(ck['model_state_dict']) |
| previous.load_state_dict(ck['previous_state_dict']) |
| optimizer.load_state_dict(ck['optimizer_state_dict']) |
| completed, tau, weights = ck['step'], ck['tau'], ck['loss_weights'] |
| initial_loss = ck['initial_residual_loss'] |
| best_score, best_step = ck['best_physics_score'], ck['best_step'] |
| previous_seconds = ck['elapsed_seconds'] |
| torch.set_rng_state(ck['torch_rng'].cpu()) |
| torch.cuda.set_rng_state_all([state.cpu() for state in ck['cuda_rng']]) |
| np.random.set_state(ck['numpy_rng']) |
| random.setstate(ck['python_rng']) |
| atomic_json(results / 'run_config.json', cfg) |
| source_files = list((root / 'model').glob('*.py')) + list((root / 'scripts').glob('*.py')) + [Path(args.config).resolve()] |
| atomic_json(results / 'source_manifest.json', { |
| 'files_sha256': {str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest() for p in source_files}, |
| 'python': sys.version, 'torch': str(torch.__version__), 'hip': torch.version.hip, |
| 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), 'host': socket.gethostname()}) |
| print(json.dumps({'event': 'training_started', 'requested_steps': requested, |
| 'resumed_at': completed, 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), |
| 'host': socket.gethostname(), 'device': torch.cuda.get_device_name(), |
| 'torch': torch.__version__, 'batch_size': cfg['training']['batch_size'], |
| 'parameters': sum(p.numel() for p in parameters), |
| 'reference_training_labels': False}), flush=True) |
| fields = ['step', 'elapsed_seconds', 'learning_rate', 'loss'] + list(LOSS_NAMES) + [ |
| 'raw_u_res', 'raw_v_res', 'raw_p_res', 'tau_u', 'tau_v', 'tau_p', |
| 'gamma_u', 'gamma_v', 'gamma_p', 'relative_l2_u', 'relative_l2_v', 'physics_validation'] |
| history_path = results / 'training_history.csv' |
| hist = history_path.open('a' if completed else 'w', newline='', buffering=1) |
| writer = csv.DictWriter(hist, fieldnames=fields) |
| if not completed: |
| writer.writeheader() |
|
|
| def checkpoint(): |
| return {'schema_version': 1, 'model_state_dict': model.state_dict(), |
| 'previous_state_dict': previous.state_dict(), |
| 'optimizer_state_dict': optimizer.state_dict(), 'config': cfg, |
| 'step': completed, 'requested_steps': requested, |
| 'tau': tau, 'loss_weights': weights, 'initial_residual_loss': initial_loss, |
| 'best_physics_score': best_score, 'best_step': best_step, |
| 'elapsed_seconds': previous_seconds + time.monotonic() - start, |
| 'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(), |
| 'numpy_rng': np.random.get_state(), 'python_rng': random.getstate(), |
| 'selection': 'fixed independent unlabeled PDE+boundary validation, no reference labels', |
| 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), 'host': socket.gethostname()} |
|
|
| last_reference_metrics = {} |
| try: |
| for index in range(completed, requested): |
| k = index + 1 |
| xy = sample_interior(cfg, device) |
| optimizer.zero_grad(set_to_none=True) |
| value, raw_residual = model.physics(xy, nu) |
| with torch.no_grad(): |
| previous_value = previous(xy) |
| if initial_loss is None: |
| initial_loss = raw_residual.detach().square().mean(0) |
| if k % cfg['pseudo_time']['update_every'] == 0: |
| _, previous_residual = previous.physics(xy, nu) |
| tau, gamma = update_tau(tau, value.detach(), previous_value, |
| raw_residual.detach(), previous_residual, initial_loss, cfg) |
| relaxed = raw_residual + (value - previous_value) / tau |
| losses = torch.cat((boundary_losses(model, boundary), relaxed.square().mean(0))) |
| if k % cfg['loss_weighting']['update_every'] == 0: |
| weights, norms = gradnorm_weights(losses, parameters, weights, cfg) |
| print(json.dumps({'event': 'adaptive_update', 'step': k, 'tau': tau.tolist(), |
| 'gamma': gamma.tolist(), 'loss_weights': weights.tolist(), |
| 'gradient_norms': norms.tolist()}), flush=True) |
| loss = (weights * losses).sum() |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f'Nonfinite training loss before update {k}') |
| loss.backward() |
| if not torch.stack([torch.isfinite(p.grad).all() for p in parameters if p.grad is not None]).all(): |
| raise FloatingPointError(f'Nonfinite parameter gradient before update {k}') |
| |
| previous.load_state_dict(model.state_dict()) |
| lr = learning_rate(index, cfg) |
| for group in optimizer.param_groups: |
| group['lr'] = lr |
| optimizer.step() |
| completed = k |
| report = k == 1 or k % cfg['training']['log_every'] == 0 or k == requested |
| evaluate = k == 1 or k % cfg['training']['evaluate_every'] == 0 or k == requested |
| row = {'step': k, 'elapsed_seconds': previous_seconds + time.monotonic() - start, |
| 'learning_rate': lr, 'loss': float(loss.detach())} |
| if report or evaluate: |
| row.update(zip(LOSS_NAMES, losses.detach().tolist())) |
| row.update(zip(('raw_u_res', 'raw_v_res', 'raw_p_res'), raw_residual.detach().square().mean(0).tolist())) |
| row.update(zip(('tau_u', 'tau_v', 'tau_p'), tau.tolist())) |
| row.update(zip(('gamma_u', 'gamma_v', 'gamma_p'), gamma.tolist())) |
| if evaluate: |
| pred = predict(model, reference_xy, device, cfg['evaluation']['batch_size']) |
| last_reference_metrics = errors(pred, reference_uvp) |
| score = physical_validation(model, validation_xy, boundary, nu) |
| row.update(last_reference_metrics) |
| row['physics_validation'] = score |
| if score < best_score: |
| best_score, best_step = score, k |
| atomic_save(root / cfg['paths']['best_checkpoint'], checkpoint()) |
| atomic_json(results / 'progress.json', dict(status='running', completed_steps=k, |
| requested_steps=requested, best_step=best_step, best_physics_score=best_score, |
| elapsed_seconds=row['elapsed_seconds'], **last_reference_metrics)) |
| if k % cfg['training']['checkpoint_every'] == 0 or k == requested: |
| atomic_save(root / cfg['paths']['last_checkpoint'], checkpoint()) |
| if report or evaluate: |
| writer.writerow(row) |
| print(json.dumps({'event': 'train', **row}), flush=True) |
| summary = {'status': 'completed', 'requested_steps': requested, 'completed_steps': completed, |
| 'best_step': best_step, 'best_physics_score': best_score, |
| 'elapsed_seconds': previous_seconds + time.monotonic() - start, |
| 'final_metrics': last_reference_metrics, 'paper_relative_l2_u': cfg['evaluation']['paper_relative_l2_u'], |
| 'early_stopping': False, 'reference_used_for_optimization': False, |
| 'slurm_job_id': os.environ.get('SLURM_JOB_ID'), 'host': socket.gethostname()} |
| atomic_json(results / 'training_summary.json', summary) |
| atomic_json(results / 'progress.json', summary) |
| print(json.dumps({'event': 'training_completed', **summary}), flush=True) |
| except BaseException as exc: |
| detail = {'status': 'failed', 'completed_steps': completed, 'requested_steps': requested, |
| 'error_type': type(exc).__name__, 'error': str(exc), 'traceback': traceback.format_exc(), |
| 'elapsed_seconds': previous_seconds + time.monotonic() - start, |
| 'slurm_job_id': os.environ.get('SLURM_JOB_ID')} |
| atomic_json(results / 'failure.json', detail) |
| atomic_json(results / 'progress.json', detail) |
| if completed: |
| atomic_save(root / 'weight/interrupted_model.pth', checkpoint()) |
| print(json.dumps(detail), flush=True) |
| raise |
| finally: |
| hist.close() |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
