scripts/train.py

#!/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()