#!/usr/bin/env python3 """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) # Equation 2.62 acts on tau itself. No inverse-tau shrink or undocumented clip. 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'] # Equations C.17--C.20 specify a SUM, rather than the repository's mean. 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 is held in this evaluation-only branch; it never enters losses. 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}') # Copy theta^k BEFORE applying the optimizer: theta^(k-1) for next step. 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()