| """Validation-only model selection with restartable optimizer/RNG snapshots.""" |
| import os |
| os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib')) |
| import argparse,json,time,random,datetime,copy |
| from pathlib import Path |
| from recovery_state import completed, atomic_json |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader,TensorDataset |
| from braindecode.models import EEGNet,EEGConformer |
| from bci_data import cached,preprocess,montage |
| from classical import metrics |
|
|
| def build(kind,n,names=None): |
| if kind.startswith(('reve_','csbrain_')): |
| from foundation_models import FoundationClassifier |
| if names is None: |
| _,_,names=cached(1,'train');idx=montage(names,n);names=[names[i] for i in idx] |
| base,mode=kind.split('_',1);return FoundationClassifier(base,names,mode) |
| if kind.startswith('cbramod_'): |
| from cbramod_probe import backbone |
| m=nn.Sequential(backbone(),nn.Flatten(1),nn.LayerNorm(n*2*200),nn.Dropout(.5),nn.Linear(n*2*200,5)) |
| if kind=='cbramod_partial': |
| for p in m[0].parameters():p.requires_grad=False |
| for p in m[0].encoder.layers[-2:].parameters():p.requires_grad=True |
| return m |
| if kind=='eegnet':return EEGNet(n_chans=n,n_outputs=5,n_times=400,F1=8,D=2,kernel_length=64,drop_prob=.5) |
| if kind=='conformer':return EEGConformer(n_chans=n,n_outputs=5,n_times=400,n_filters_time=40,att_depth=3,att_heads=10,drop_prob=.5,att_drop_prob=.5) |
| raise ValueError(kind) |
|
|
| def expired(): |
| from gpu_utils import STOP |
| if STOP:return True |
| from recovery_state import deadline |
| return datetime.datetime.now(datetime.timezone.utc)>=deadline() |
|
|
| def run(kind,n,subjects,seed,epochs=80,lr=.001,selection='distributed'): |
| tag='pooled' if len(subjects)>1 else f'sub-{subjects[0]:02d}' |
| name=f'strict_{kind}_{selection}_{n}_{tag}_seed-{seed}';out=Path('experiments')/(name+'.json') |
| if completed(name):return |
| torch.set_num_threads(4);random.seed(seed);np.random.seed(seed);torch.manual_seed(seed) |
| device='cuda' if torch.cuda.is_available() else 'cpu';t=time.time() |
| if os.environ.get('SETU_REQUIRE_CUDA')=='1': |
| from gpu_utils import require_cuda |
| require_cuda() |
| use_amp=device=='cuda' and os.environ.get('SETU_GPU_AMP')=='1' |
| xx=[];yy=[];vv=[];zz=[] |
| for s in subjects: |
| x,y,names=cached(s,'train');v,z,_=cached(s,'validation');idx=montage(names,n,selection) |
| xx.append(preprocess(x,idx));yy.append(y);vv.append(preprocess(v,idx));zz.append(z) |
| x=np.concatenate(xx);y=np.concatenate(yy);v=np.concatenate(vv);z=np.concatenate(zz) |
| scale=np.full((1,n,1),100.,dtype=np.float32) if kind.startswith('cbramod_') else np.maximum(x.std(axis=(0,2),keepdims=True),1e-3) |
| x=np.clip(x/scale,-20,20);v=np.clip(v/scale,-20,20) |
| if kind.startswith('cbramod_'):x=x.reshape(-1,n,2,200);v=v.reshape(-1,n,2,200) |
| loader=DataLoader(TensorDataset(torch.tensor(x),torch.tensor(y)),batch_size=16 if kind.startswith('cbramod_') else 64,shuffle=True,num_workers=0) |
| vx=torch.tensor(v,device=device);model=build(kind,n).to(device);opt=torch.optim.AdamW([{'params':model[0].parameters(),'lr':lr/10},{'params':list(model[2].parameters())+list(model[4].parameters()),'lr':lr}] if kind.startswith('cbramod_') else model.parameters(),lr=lr,weight_decay=.01) |
| weights=len(y)/(5*np.bincount(y,minlength=5));criterion=nn.CrossEntropyLoss(weight=torch.tensor(weights,dtype=torch.float32,device=device),label_smoothing=.05) |
| best=-1;best_epoch=0;history=[];start=0;resume_device=None |
| last=Path('checkpoints')/(name+'_last.pt');bestpath=Path('checkpoints')/(name+'.pt') |
| if last.exists(): |
| state=torch.load(last,map_location='cpu',weights_only=False);resume_device='cuda' if state.get('cuda_rng') is not None else 'cpu';model.load_state_dict(state['model']);opt.load_state_dict(state['optimizer']);start=state['epoch']+1;best=state['best'];best_epoch=state['best_epoch'];history=state['history'];torch.set_rng_state(state['rng']) |
| if device=='cuda' and state.get('cuda_rng') is not None:torch.cuda.set_rng_state_all(state['cuda_rng']) |
| if 'python_rng' in state:random.setstate(state['python_rng']) |
| if 'numpy_rng' in state:np.random.set_state(state['numpy_rng']) |
|
|
| if not last.exists() and bestpath.exists(): |
| state=torch.load(bestpath,map_location='cpu',weights_only=False);model.load_state_dict(state['model']);best=state['validation_metrics']['balanced_accuracy'];best_epoch=state['epoch'];start=best_epoch+1 |
| print(name,'Resuming best weights only: optimizer/RNG snapshot unavailable',flush=True) |
| if last.exists() and start-1-best_epoch>=15:start=epochs |
| for epoch in range(start,epochs): |
| if expired():print('WALL CLOCK LIMIT',flush=True);return |
| model.train();losses=[] |
| for a,b in loader: |
| a=a.to(device);b=b.to(device) |
| a=a*(.9+.2*torch.rand((len(a),)+(1,)*(a.ndim-1),device=device))+torch.randn_like(a)*.02 |
| opt.zero_grad(set_to_none=True) |
| with torch.autocast(device_type=device,dtype=torch.bfloat16,enabled=use_amp):logits=model(a);loss=criterion(logits,b) |
| if not torch.isfinite(loss):raise FloatingPointError('Nonfinite training loss') |
| loss.backward();nn.utils.clip_grad_norm_(model.parameters(),1);opt.step();losses.append(loss.item()) |
| model.eval() |
| with torch.no_grad(),torch.autocast(device_type=device,dtype=torch.bfloat16,enabled=use_amp):pred=torch.cat([model(q) for q in vx.split(16 if kind.startswith('cbramod_') else 64)]).argmax(1).cpu().numpy() |
| m=metrics(z,pred);score=m['balanced_accuracy'];history.append({'epoch':epoch,'loss':float(np.mean(losses)),**m}) |
| if score>best: |
| best=score;best_epoch=epoch |
| from gpu_utils import atomic_save |
| atomic_save({'model':model.cpu().state_dict(),'kind':kind,'channels':n,'indices':idx,'names':names,'scale':scale,'seed':seed,'subjects':subjects,'epoch':epoch,'validation_metrics':m,'training_device':device,'precision':'bfloat16' if use_amp else 'float32'},bestpath);model.to(device) |
| from gpu_utils import atomic_save |
| atomic_save({'model':model.state_dict(),'optimizer':opt.state_dict(),'epoch':epoch,'best':best,'best_epoch':best_epoch,'history':history,'rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all() if device=='cuda' else None,'python_rng':random.getstate(),'numpy_rng':np.random.get_state()},last) |
| print(name,epoch,round(score,4),round(float(np.mean(losses)),4),flush=True) |
| if epoch-best_epoch>=15:break |
| if expired():return |
| best_state=torch.load(bestpath,map_location='cpu',weights_only=False);model.load_state_dict(best_state['model']);model.eval() |
| with torch.inference_mode(),torch.autocast(device_type=device,dtype=torch.bfloat16,enabled=use_amp): final_pred=torch.cat([model(q) for q in vx.split(16 if kind.startswith('cbramod_') else 64)]).argmax(1).cpu().numpy() |
| Path('reports/validation_predictions').mkdir(exist_ok=True) |
| np.savez_compressed(f'reports/validation_predictions/{name}.npz',y=z,pred=final_pred) |
| subject_metrics={str(subject):metrics(z[i*50:(i+1)*50],final_pred[i*50:(i+1)*50]) for i,subject in enumerate(subjects)} |
| atomic_json(out, {'name':name,'model':kind,'channels':n,'montage':selection,'subjects':subjects,'subject_validation_metrics':subject_metrics,'seed':seed,'split':'validation','best_epoch':best_epoch,'metrics':history[best_epoch]['balanced_accuracy'] if False else torch.load(bestpath,weights_only=False)['validation_metrics'],'history':history,'parameters':sum(p.numel() for p in model.parameters()),'seconds':time.time()-t,'device':device,'resumed_checkpoint_device':resume_device,'best_checkpoint_training_device':best_state.get('training_device','cpu'),'precision':'bfloat16' if use_amp else 'float32','lr':lr}) |
| if __name__=='__main__': |
| from experiment_guard import assert_tuning_open |
| assert_tuning_open() |
| if os.environ.get('SETU_REQUIRE_CUDA')=='1': |
| from gpu_utils import require_cuda,install_signals |
| require_cuda();install_signals() |
| p=argparse.ArgumentParser();p.add_argument('--model',default='eegnet');p.add_argument('--channels',nargs='+',type=int,default=[64]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));p.add_argument('--seeds',nargs='+',type=int,default=[2026]);p.add_argument('--epochs',type=int,default=80);p.add_argument('--pooled',action='store_true');p.add_argument('--montage',default='distributed');p.add_argument('--lr',type=float,default=.001);a=p.parse_args() |
| for n in a.channels: |
| for seed in a.seeds: |
| for ss in ([a.subjects] if a.pooled else [[s] for s in a.subjects]):run(a.model,n,ss,seed,a.epochs,a.lr,a.montage) |
|
|
