File size: 7,552 Bytes
4fcc185 | 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 | """CUDA-only foundation adaptation. Official train/validation only; resumable epochs."""
import os,json,time,argparse,random
from pathlib import Path
from recovery_state import completed, atomic_json
import numpy as np,torch,joblib
from torch import nn
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from bci_data import cached,preprocess,montage
from classical import metrics
from foundation_models import Foundation,FoundationClassifier
from gpu_utils import require_cuda,expired,atomic_save,cpu_state,install_signals
def scale_for(kind):return 1000. if kind=='reve' else 100.
def tensors(subject,n,kind):
x,y,names=cached(subject,'train');v,z,_=cached(subject,'validation');idx=montage(names,n);scale=scale_for(kind)
return preprocess(x,idx)/scale,y,preprocess(v,idx)/scale,z,idx,names
def probe(kind,n,s):
name=f'strict_{kind}_probe_distributed_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json')
if completed(name):return
device=require_cuda();x,y,v,z,idx,names=tensors(s,n,kind);encoder=Foundation(kind,[names[i] for i in idx]).to(device).eval();features=[];start=time.time()
for split,a in [('train',x),('validation',v)]:
path=Path('preprocessing')/f'{name}_{split}_features.npy'
if path.exists():features.append(np.load(path));continue
chunks=[]
with torch.inference_mode():
for q in torch.tensor(a).split(16):
if expired():return
with torch.autocast('cuda',dtype=torch.bfloat16):f=encoder(q.to(device))
chunks.append(f.float().cpu().numpy())
f=np.concatenate(chunks);np.save(path,f);features.append(f)
del encoder;torch.cuda.empty_cache();best=None;search=[]
for c in [.001,.01,.1,1.]:
if expired():return
model=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1500,class_weight='balanced'));model.fit(features[0],y);m=metrics(z,model.predict(features[1]));search.append({'C':c,**m})
if best is None or m['balanced_accuracy']>best[0]:best=(m['balanced_accuracy'],model,c,m)
Path('reports/validation_predictions').mkdir(exist_ok=True)
np.savez_compressed(f'reports/validation_predictions/{name}.npz',y=z,pred=best[1].predict(features[1]))
joblib.dump({'model':best[1],'kind':kind,'indices':idx,'names':names,'subject':s,'scale':scale_for(kind),'preprocessing':'car_1-40_imagined0to2s_200hz','feature_precision':'cuda_bfloat16'},f'checkpoints/{name}.joblib')
atomic_json(out, {'name':name,'model':kind+'_probe','subject':s,'channels':n,'montage':'distributed','split':'validation','selected_C':best[2],'metrics':best[3],'search':search,'seconds':time.time()-start,'device':'cuda','precision':'bfloat16'});print(name,best[3],flush=True)
def finetune(kind,mode,n,s,epochs=50,lr=.0001,seed=2026):
name=f'strict_{kind}_{mode}_distributed_{n}_sub-{s:02d}_seed-{seed}';out=Path('experiments')/(name+'.json')
if completed(name):return
device=require_cuda();random.seed(seed);np.random.seed(seed);torch.manual_seed(seed);torch.cuda.manual_seed_all(seed);starttime=time.time()
x,y,v,z,idx,names=tensors(s,n,kind);x=np.clip(x,-20,20);v=np.clip(v,-20,20);model=FoundationClassifier(kind,[names[i] for i in idx],mode).to(device)
opt=torch.optim.AdamW([{'params':[p for p in model.encoder.parameters() if p.requires_grad],'lr':lr/10},{'params':model.head.parameters(),'lr':lr}],weight_decay=.01)
loader=torch.utils.data.DataLoader(torch.utils.data.TensorDataset(torch.tensor(x),torch.tensor(y)),batch_size=16,shuffle=True);vx=torch.tensor(v,device=device)
best=-1;best_epoch=-1;start=0;history=[];last=Path(f'checkpoints/{name}_last.pt');bestpath=Path(f'checkpoints/{name}.pt')
if last.exists():
state=torch.load(last,map_location='cpu',weights_only=False);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']);torch.cuda.set_rng_state_all(state['cuda_rng']);random.setstate(state['python_rng']);np.random.set_state(state['numpy_rng'])
criterion=nn.CrossEntropyLoss(label_smoothing=.05)
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>=12:start=epochs
for epoch in range(start,epochs):
if expired():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,1),device=device))+torch.randn_like(a)*(.02 if kind=='csbrain' else .002)
opt.zero_grad(set_to_none=True)
with torch.autocast('cuda',dtype=torch.bfloat16):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.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):pred=torch.cat([model(q) for q in vx.split(16)]).float().argmax(1).cpu().numpy()
m=metrics(z,pred);history.append({'epoch':epoch,'loss':float(np.mean(losses)),**m});score=m['balanced_accuracy']
if score>best:
best=score;best_epoch=epoch;atomic_save({'model':cpu_state(model),'kind':kind+'_'+mode,'channels':n,'indices':idx,'names':names,'scale':np.full((1,n,1),scale_for(kind),dtype=np.float32),'seed':seed,'subjects':[s],'epoch':epoch,'validation_metrics':m,'training_device':'cuda','precision':'bfloat16'},bestpath)
atomic_save({'model':cpu_state(model),'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(),'python_rng':random.getstate(),'numpy_rng':np.random.get_state()},last)
print(name,epoch,score,float(np.mean(losses)),flush=True)
if epoch-best_epoch>=12:break
if expired():return
beststate=torch.load(bestpath,map_location='cpu',weights_only=False)
model.load_state_dict(beststate['model']);model.eval()
with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):pred=torch.cat([model(q) for q in vx.split(16)]).float().argmax(1).cpu().numpy()
Path('reports/validation_predictions').mkdir(exist_ok=True)
np.savez_compressed(f'reports/validation_predictions/{name}.npz',y=z,pred=pred)
atomic_json(out, {'name':name,'model':kind+'_'+mode,'channels':n,'montage':'distributed','subjects':[s],'seed':seed,'split':'validation','best_epoch':best_epoch,'metrics':beststate['validation_metrics'],'history':history,'seconds':time.time()-starttime,'device':'cuda','precision':'bfloat16','lr':lr,'parameters':sum(p.numel() for p in model.parameters())})
if __name__=='__main__':
from experiment_guard import assert_tuning_open
assert_tuning_open();require_cuda();install_signals();torch.set_num_threads(4)
p=argparse.ArgumentParser();p.add_argument('--model',choices=['reve','csbrain'],required=True);p.add_argument('--mode',choices=['probe','partial','full'],default='probe');p.add_argument('--channels',nargs='+',type=int,default=[64,32,16,8]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));p.add_argument('--epochs',type=int,default=50);a=p.parse_args()
for n in a.channels:
for s in a.subjects:
if expired():raise SystemExit(3)
probe(a.model,n,s) if a.mode=='probe' else finetune(a.model,a.mode,n,s,a.epochs)
|