#!/usr/bin/env python """Formal 100000-step online physics training. No shortened-run mode.""" import argparse import csv import json import math import os from pathlib import Path import random import signal 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.comphy_tgv import sample_training_points from scripts.common import (ROOT, resolve, load_config, build_model, load_reference, metrics, predict, atomic_checkpoint, atomic_json, sha256) class Tee: def __init__(self, original, file): self.original, self.file = original, file def write(self, message): self.original.write(message) self.file.write(message) self.file.flush() def flush(self): self.original.flush() self.file.flush() def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--config', default=str(ROOT/'config/config.yaml')) parser.add_argument('--resume', type=Path) args = parser.parse_args() cfg = load_config(args.config) tr = cfg['training'] if not os.environ.get('SLURM_JOB_ID'): raise RuntimeError('Formal training must run on an allocated SLURM compute node.') device = torch.device(tr['device']) if device.type != 'cuda' or not torch.cuda.is_available(): raise RuntimeError('Formal training requires an accessible allocated accelerator.') torch.set_num_threads(1) torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False random.seed(tr['seed']); np.random.seed(tr['seed']); torch.manual_seed(tr['seed']) torch.cuda.manual_seed_all(tr['seed']) results, weights = resolve(cfg['paths']['results']), resolve(cfg['paths']['weight']) results.mkdir(parents=True, exist_ok=True); weights.mkdir(parents=True, exist_ok=True) log = open(results/'train.log', 'a', buffering=1) sys.stdout = Tee(sys.stdout, log) model = build_model(cfg, device) optimizer = torch.optim.Adam(model.parameters(), lr=tr['learning_rate'], betas=(0.9, 0.999), eps=1e-8, weight_decay=0.) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, tr['decay_steps'], gamma=tr['decay_factor']) # These arrays can only enter no_grad evaluation, never loss_terms/train_step. coordinates, target = load_reference(cfg['validation']['reference']) step, best, best_step, previous_seconds = 0, math.inf, 0, 0. if args.resume: state = torch.load(args.resume, map_location=device, weights_only=False) if state['config'] != cfg: raise ValueError('Resume config differs from original trajectory.') model.load_state_dict(state['model_state_dict']) optimizer.load_state_dict(state['optimizer_state_dict']) scheduler.load_state_dict(state['scheduler_state_dict']) step, best, best_step = state['step'], state['best_l2_error'], state['best_step'] previous_seconds = state['elapsed_seconds'] torch.set_rng_state(state['torch_rng'].cpu()) torch.cuda.set_rng_state_all([s.cpu() for s in state['cuda_rng']]) np.random.set_state(state['numpy_rng']); random.setstate(state['python_rng']) elif (weights/'last_model.pth').exists(): raise FileExistsError('Existing training state; use --resume weight/last_model.pth.') info = {'job_id':os.environ['SLURM_JOB_ID'], 'hostname':socket.gethostname(), 'python':sys.version, 'torch':torch.__version__, 'hip':torch.version.hip, 'device':torch.cuda.get_device_name(device), 'device_properties':str(torch.cuda.get_device_properties(device)), 'config':cfg, 'parameters':sum(p.numel() for p in model.parameters()), 'reference_sha256':sha256(cfg['validation']['reference']), 'reference_usage':'validation/selection only; no_grad; not in optimization loss'} atomic_json(results/'environment.json',info) print('FORMAL_RUN '+json.dumps(info),flush=True) print('SAMPLING interior=1000 initial=1000 boundary=1000 freshly independent each step; reference validation only',flush=True) stopped = [] def stop_handler(signum, frame): stopped.append(signum) for sig in (signal.SIGTERM, signal.SIGINT, signal.SIGUSR1): signal.signal(sig, stop_handler) start = time.monotonic() def elapsed(): return previous_seconds + time.monotonic()-start def checkpoint(): return {'model_state_dict':model.state_dict(), 'optimizer_state_dict':optimizer.state_dict(), 'scheduler_state_dict':scheduler.state_dict(), 'step':step, 'best_l2_error':best, 'best_step':best_step, 'elapsed_seconds':elapsed(), 'config':cfg, 'torch_rng':torch.get_rng_state(), 'cuda_rng':torch.cuda.get_rng_state_all(), 'numpy_rng':np.random.get_state(), 'python_rng':random.getstate(), 'reference_sha256':info['reference_sha256']} def evaluate(): nonlocal best, best_step measurement = metrics(predict(model,coordinates,cfg['validation']['batch_size']),target) if measurement['l2_error'] < best: best, best_step = measurement['l2_error'], step payload = checkpoint(); payload['validation_metrics'] = measurement atomic_checkpoint(weights/'best_model.pth',payload) record = {'step':step,'elapsed_seconds':elapsed(), **measurement, 'best_l2_error':best, 'best_step':best_step} with open(results/'validation_history.jsonl','a') as f: f.write(json.dumps(record)+'\n') print('VALIDATION '+json.dumps(record),flush=True) return measurement keys = list(tr['loss_weights']) history = open(results/'training_history.csv','a',newline='',buffering=1) writer = csv.DictWriter(history, fieldnames=['step','elapsed_seconds','learning_rate','loss']+keys) if history.tell() == 0: writer.writeheader() status, reason = 'running', None last_evaluation = None try: if step == 0: last_evaluation = evaluate() for current_step in range(step+1, tr['steps']+1): model.train() batch = sample_training_points(tr['interior_points'],tr['initial_points'], tr['boundary_points'],device,next(model.parameters()).dtype) optimizer.zero_grad(set_to_none=True) terms = model.loss_terms(*batch) loss = sum(terms[k]*tr['loss_weights'][k] for k in keys) if not torch.isfinite(loss).item(): raise FloatingPointError(f'Nonfinite loss before update {current_step}') loss.backward() lr_used = optimizer.param_groups[0]['lr'] optimizer.step(); scheduler.step() step = current_step if step == 1 or step % tr['log_every'] == 0: row = {'step':step,'elapsed_seconds':elapsed(),'learning_rate':lr_used, 'loss':float(loss.detach()),**{k:float(v.detach()) for k,v in terms.items()}} writer.writerow(row) print('TRAIN '+json.dumps(row),flush=True) atomic_json(results/'training_status.json', {'status':'running','actual_steps':step, 'requested_steps':tr['steps'],'elapsed_seconds':elapsed(), 'best_step':best_step,'best_l2_error':best,'reason':None}) if step % tr['validate_every'] == 0 or step == tr['steps']: last_evaluation = evaluate() if step % tr['checkpoint_every'] == 0 or stopped: atomic_checkpoint(weights/'last_model.pth',checkpoint()) if stopped: status, reason = 'interrupted', f'Received signal {stopped[-1]}' break else: status = 'completed' except BaseException as exc: status, reason = 'failed', repr(exc) traceback.print_exc() raise finally: if math.isfinite(best): atomic_checkpoint(weights/'last_model.pth',checkpoint()) summary={'status':status,'requested_steps':tr['steps'],'actual_steps':step, 'best_step':best_step,'best_l2_error':best if math.isfinite(best) else None, 'last_validation':last_evaluation,'elapsed_seconds':elapsed(),'reason':reason, 'full_training_completed':step==tr['steps'] and status=='completed'} atomic_json(results/'training_status.json',summary) print('TRAINING_FINAL '+json.dumps(summary),flush=True) history.close() if status != 'completed': sys.exit(2) if __name__ == '__main__': main()