scripts/train.py

"""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'])
    # One graph for the IC and the original + partner BC coordinates.
    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:]
    # Sum of 12 per-constraint MSEs (four sides, three variables).
    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()