| |
| """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']) |
| |
| 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() |
|
|
