File size: 9,451 Bytes
3c8d50e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """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):
# Eq. 44 is the function marginal; do not apply a second exponential.
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)
|