| import os,sys,time,json,hashlib,argparse |
| from pathlib import Path |
| os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib')) |
| sys.path.insert(0,os.path.abspath('external/CBraMod')) |
| import torch,numpy as np,joblib |
| from torch import nn |
| from bci_data import cached,preprocess,montage |
| from classical import metrics |
| from sklearn.pipeline import make_pipeline |
| from sklearn.preprocessing import StandardScaler |
| from sklearn.linear_model import LogisticRegression |
| from models.cbramod import CBraMod |
|
|
| def backbone(): |
| m=CBraMod();state=torch.load('models/cbramod_pretrained.pth',map_location='cpu',weights_only=True);m.load_state_dict(state,strict=True);m.proj_out=nn.Identity();return m |
|
|
| def get_features(m,subject,split,n,device,evaluation=False): |
| p=Path('preprocessing')/f'strict_cbramod_{n}_sub-{subject:02d}_{split}.npy' |
| x,y,names=cached(subject,split,evaluation) |
| idx=montage(names,n) |
| if p.exists():return np.load(p),y,idx |
| x=preprocess(x,idx)/100 |
| feats=[];m.eval() |
| with torch.inference_mode(): |
| for a in torch.tensor(x).split(16):feats.append(m(a.to(device).reshape(-1,n,2,200)).flatten(1).cpu().numpy()) |
| f=np.concatenate(feats);np.save(p,f);return f,y,idx |
| if __name__=='__main__': |
| from experiment_guard import assert_tuning_open |
| assert_tuning_open() |
| p=argparse.ArgumentParser();p.add_argument('--channels',nargs='+',type=int,default=[64]);p.add_argument('--subjects',nargs='+',type=int,default=list(range(1,16)));a=p.parse_args() |
| torch.set_num_threads(4);device='cuda' if torch.cuda.is_available() else 'cpu';m=backbone().to(device) |
| print('strict official checkpoint loaded',sum(p.numel() for p in m.parameters()),device,flush=True) |
| for n in a.channels: |
| for s in a.subjects: |
| name=f'strict_cbramod_distributed_{n}_sub-{s:02d}';out=Path('experiments')/(name+'.json') |
| if out.exists():continue |
| from neural import expired |
| if expired():sys.exit(0) |
| t=time.time();x,y,idx=get_features(m,s,'train',n,device);v,z,_=get_features(m,s,'validation',n,device);best=None;search=[] |
| for c in [.0001,.001,.01,.1,1]: |
| clf=make_pipeline(StandardScaler(),LogisticRegression(C=c,max_iter=1000,class_weight='balanced'));clf.fit(x,y);met=metrics(z,clf.predict(v));search.append({'C':c,**met}) |
| if best is None or met['balanced_accuracy']>best[0]:best=(met['balanced_accuracy'],clf,c,met) |
| joblib.dump({'model':best[1],'kind':'cbramod','indices':idx,'subject':s,'preprocessing':'strict_0to2s_uVdiv100','checkpoint_sha256':hashlib.sha256(Path('models/cbramod_pretrained.pth').read_bytes()).hexdigest()},Path('checkpoints')/(name+'.joblib')) |
| out.write_text(json.dumps({'name':name,'model':'cbramod','channels':n,'montage':'distributed','subject':s,'split':'validation','metrics':best[3],'C':best[2],'search':search,'seconds':time.time()-t,'device':device,'mode':'frozen_flattened_features_linear_probe'},indent=2));print(name,best[0],time.time()-t,flush=True) |
|
|
