| """Train the paper-budget DeepONet with hybrid exponential vRBA.""" |
| import csv |
| import json |
| import math |
| import random |
| import sys |
| import time |
| from pathlib import Path |
| sys.path.insert(0,str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import torch |
| from model.vrba import VRBADeepONet, ExponentialVRBA |
| from common import arguments, config, path, accelerator, load_data, write_json, save_checkpoint, dataset_manifest |
|
|
|
|
| def run(cfg, resume=None, smoke=False): |
| seed=cfg['training']['seed'] |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) |
| torch.set_num_threads(1) |
| device=accelerator(cfg['training']['device']) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cuda.matmul.allow_tf32=False |
| torch.backends.cudnn.allow_tf32=False |
| data=load_data(cfg) |
| results=path(cfg['paths']['results']) / 'smoke' if smoke else path(cfg['paths']['results']) |
| results.mkdir(parents=True,exist_ok=True) |
| if not smoke and (results/'history.csv').exists() and resume is None: |
| raise RuntimeError('Existing formal history; use --resume or a new results directory.') |
| write_json(results/'data_manifest.json',dataset_manifest(data,cfg)) |
| write_json(results/'config_resolved.json',cfg) |
| x=torch.as_tensor(data['x'],device=device) |
| y=torch.as_tensor(data['y'],device=device) |
| t=torch.as_tensor(data['time'],device=device) |
| nt=cfg['data']['train']; bs=cfg['training']['batch_size'] |
| steps=nt//bs |
| if nt%bs: raise ValueError('Training count must be divisible by batch size.') |
| epochs=1 if smoke else cfg['training']['epochs'] |
| model=VRBADeepONet(**cfg['model']).to(device) |
| nparam=sum(p.numel() for p in model.parameters()) |
| if nparam != 101100: raise ValueError(f'Paper parameter count mismatch: {nparam}') |
| v=cfg['vrba'] |
| attention=ExponentialVRBA(nt,cfg['data']['output_points'],device, |
| **{k:v[k] for k in ['eta','gamma','phi','c','initial_weight']}) |
| optimizer=torch.optim.Adam(model.parameters(),lr=cfg['training']['learning_rate'], |
| weight_decay=cfg['training']['weight_decay'],betas=tuple(cfg['training']['adam_betas']), |
| eps=cfg['training']['adam_epsilon']) |
| scheduler=torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max=cfg['training']['epochs']*steps) |
| best=float('inf'); best_epoch=0; start_epoch=0; step=0 |
| if resume: |
| ck=torch.load(path(resume),map_location=device,weights_only=False) |
| if ck['model_config']!=cfg['model']: raise ValueError('Resume architecture mismatch.') |
| if ck['normalization']!=data['stats']: raise ValueError('Resume normalization mismatch.') |
| model.load_state_dict(ck['model_state_dict']); optimizer.load_state_dict(ck['optimizer_state_dict']) |
| scheduler.load_state_dict(ck['scheduler_state_dict']); attention.weights.copy_(ck['attention_weights'].to(device)) |
| start_epoch=ck['epoch']; step=ck['global_step']; best=ck['best_validation']; best_epoch=ck['best_epoch'] |
| torch.set_rng_state(ck['torch_rng'].cpu()); torch.cuda.set_rng_state_all([r.cpu() for r in ck['cuda_rng']]) |
| np.random.set_state(ck['numpy_rng']); random.setstate(ck['python_rng']) |
| val=data['splits']['validation'] |
| fields=['epoch','global_step','train_weighted_loss','train_sampled_mse_normalized', |
| 'validation_mse_normalized','validation_relative_l2_normalized_batch_mean', |
| 'validation_relative_l2_physical_global','learning_rate','best_epoch','elapsed_seconds'] |
| mode='a' if resume else 'w' |
| begin=time.perf_counter() |
| environment=dict(python=sys.version,torch=torch.__version__,hip=torch.version.hip, |
| device=torch.cuda.get_device_name(device),parameter_count=nparam, |
| job_id=__import__('os').environ.get('SLURM_JOB_ID'),node=__import__('socket').gethostname()) |
| write_json(results/'environment.json',environment) |
| print('RUN_START '+json.dumps(dict(environment=environment,epochs=epochs,batch_size=bs,steps_per_epoch=steps,smoke=smoke)),flush=True) |
| checks={} |
| with (results/'history.csv').open(mode,newline='',buffering=1) as f: |
| writer=csv.DictWriter(f,fieldnames=fields) |
| if not resume: writer.writeheader() |
| for epoch in range(start_epoch,epochs): |
| model.train(); losses=torch.zeros(2,device=device) |
| for _ in range(steps): |
| |
| with torch.no_grad(): |
| indices=torch.multinomial(attention.probabilities(),bs,replacement=v['replacement']) |
| optimizer.zero_grad(set_to_none=True) |
| prediction=model(x[indices],t) |
| if smoke and not checks: |
| assert prediction.shape==(bs,cfg['data']['output_points']) |
| checks['forward_pass']='passed' |
| residual=prediction-y[indices] |
| weights=attention.update(indices,residual,step) |
| loss=(weights*residual).square().mean() |
| loss.backward() |
| if smoke and 'backward_pass' not in checks: |
| assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters()) |
| checks['backward_pass']='passed' |
| optimizer.step(); scheduler.step(); step+=1 |
| losses+=torch.stack((loss.detach(),residual.detach().square().mean())) |
| model.eval() |
| with torch.inference_mode(): |
| pv=model(x[val],t); rv=pv-y[val] |
| val_mse=rv.square().mean().item() |
| val_rel=torch.stack([torch.linalg.vector_norm(rv[i:i+bs])/torch.linalg.vector_norm(y[val][i:i+bs]) |
| for i in range(0,cfg['data']['validation'],bs)]).mean().item() |
| physical_true=y[val]*data['stats']['radius_std']+data['stats']['radius_mean'] |
| physical_rel=(torch.linalg.vector_norm(rv*data['stats']['radius_std'])/torch.linalg.vector_norm(physical_true)).item() |
| train_loss,train_mse=(losses/steps).tolist() |
| if not all(math.isfinite(vv) for vv in [train_loss,train_mse,val_mse,val_rel,physical_rel]): |
| raise FloatingPointError(f'Nonfinite values at epoch {epoch+1}') |
| improved=val_rel<best |
| if improved: best=val_rel; best_epoch=epoch+1 |
| def checkpoint(): |
| return dict(format_version=1,model_state_dict=model.state_dict(),model_config=cfg['model'], |
| normalization=data['stats'],epoch=epoch+1,global_step=step,best_validation=best,best_epoch=best_epoch, |
| validation_selection='mean of normalized relative L2 over two batches of 50', |
| optimizer_state_dict=optimizer.state_dict(),scheduler_state_dict=scheduler.state_dict(), |
| attention_weights=attention.weights.detach().cpu(),torch_rng=torch.get_rng_state(), |
| cuda_rng=torch.cuda.get_rng_state_all(),numpy_rng=np.random.get_state(),python_rng=random.getstate(), |
| data_selection='original indices 0:1000; train 0:800, val 800:900, test 900:1000', |
| training_config=cfg['training'],vrba_config=cfg['vrba']) |
| if not smoke and improved: save_checkpoint(path(cfg['paths']['best_checkpoint']),checkpoint()) |
| if not smoke and ((epoch+1)%cfg['training']['checkpoint_frequency']==0 or epoch+1==epochs): |
| save_checkpoint(path(cfg['paths']['last_checkpoint']),checkpoint()) |
| row=dict(zip(fields,[epoch+1,step,train_loss,train_mse,val_mse,val_rel,physical_rel, |
| scheduler.get_last_lr()[0],best_epoch,time.perf_counter()-begin])) |
| writer.writerow(row) |
| print('EPOCH '+json.dumps(row),flush=True) |
| if smoke: |
| domain = dict(relative_l2_normalized=val_rel, |
| residual_variance_absolute=rv.abs().var(unbiased=False).item(), |
| residual_linf=rv.abs().max().item()) |
| assert all(math.isfinite(value) for value in domain.values()) |
| write_json(results/'domain_metrics.json',domain) |
| print('SMOKE_DOMAIN_METRICS '+json.dumps(domain),flush=True) |
| checks.update(train_loop_dry_run='passed' ,validation_loop_dry_run='passed', |
| domain_metric_dry_run='passed: relative L2, residual variance and Linf; CL is not applicable to this experiment', |
| config_consistency='passed: Bx200 -> Bx100, 500x1 -> 500x100 -> Bx500; 101100 parameters') |
| assert train_loss>0 and best<1e10 |
| write_json(results/'checks.json',dict(attempt=1,checks=checks)) |
| torch.cuda.synchronize() |
| summary=dict(status='completed',completed_epochs=epochs,global_step=step,best_epoch=best_epoch, |
| best_validation_relative_l2_normalized=best,early_stopping=False,smoke=smoke, |
| duration_seconds=time.perf_counter()-begin) |
| write_json(results/'training_summary.json',summary) |
| print('TRAINING_COMPLETE '+json.dumps(summary),flush=True) |
|
|
|
|
| if __name__=='__main__': |
| parser=arguments(__doc__) |
| parser.add_argument('--resume',default=None,help='Full training-state checkpoint; restores optimizer, scheduler, vRBA and RNG.') |
| parser.add_argument('--smoke',action='store_true',help='Separate one-epoch pipeline check; no formal checkpoints.') |
| args=parser.parse_args() |
| run(config(args.config),args.resume,args.smoke) |
|
|
