| """Four preregistered small students; speaker-disjoint, causal evaluation. |
| |
| No live microphone data is read. All teachers and Nemotron remain frozen. |
| Run via python -m speaker_v3.overnight --stage train (resumes checkpoints). |
| """ |
| import html,json,os,random,time |
| from pathlib import Path |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from speaker_v2.train import roc_metrics |
|
|
| CONFIGS=[{'name':'attentive_stats','architecture':'stats','width':128}, |
| {'name':'query_gru','architecture':'gru','width':128}, |
| {'name':'causal_tcn','architecture':'tcn','width':128}, |
| {'name':'temporal_transformer','architecture':'transformer','width':128}] |
|
|
|
|
| def checkpoint(value,path): |
| temporary=path.with_suffix('.tmp.pt');torch.save(value,temporary);temporary.replace(path) |
|
|
|
|
| class CausalBlock(nn.Module): |
| def __init__(self,width,dilation): |
| super().__init__();self.d=dilation |
| self.conv=nn.Conv1d(width,width,3,dilation=dilation,groups=width) |
| self.mix=nn.Conv1d(width,width,1);self.norm=nn.LayerNorm(width) |
| def forward(self,x): |
| z=self.conv(F.pad(x.transpose(1,2),(2*self.d,0))) |
| return self.norm(x+F.gelu(self.mix(z).transpose(1,2))) |
|
|
|
|
| class Student(nn.Module): |
| """Predict newest 320-ms identity from a causal 80-ms hidden-state sequence. |
| |
| Query-conditioned attention selects relevant past frames after a change of |
| speaker. Long-window mean/std is optional, separately ablated, never required. |
| """ |
| def __init__(self,architecture='stats',width=128): |
| super().__init__();self.architecture=architecture;self.width=width |
| self.input=nn.Sequential(nn.LayerNorm(1024),nn.Linear(1024,width),nn.GELU()) |
| self.energy=nn.Linear(1,width) |
| self.long=nn.Sequential(nn.LayerNorm(2048),nn.Linear(2048,64),nn.GELU()) |
| if architecture=='gru':self.temporal=nn.GRU(width,width,batch_first=True) |
| elif architecture=='tcn':self.temporal=nn.Sequential(*[CausalBlock(width,d) for d in (1,2,4)]) |
| elif architecture=='transformer': |
| layer=nn.TransformerEncoderLayer(width,4,width*2,.1,batch_first=True,norm_first=True) |
| self.temporal=nn.TransformerEncoder(layer,2,enable_nested_tensor=False) |
| self.position=nn.Parameter(torch.randn(25,width)*.01) |
| else:self.temporal=nn.Sequential(nn.Linear(width,width),nn.GELU(),nn.Linear(width,width)) |
| self.key=nn.Linear(width,width,bias=False);self.query=nn.Linear(width,width,bias=False) |
| self.head=nn.Sequential(nn.Linear(width*3+64,384),nn.GELU(),nn.Dropout(.15)) |
| self.embedding=nn.Linear(384,192);self.activity=nn.Linear(384,3) |
| self.register_buffer('teacher_mean',torch.zeros(192));self.register_buffer('teacher_radius',torch.tensor(1.)) |
|
|
| def forward(self,x,pool,energy,use_pool=True): |
| x=self.input(x)+self.energy(((energy.float()+40)/40).clamp(-2,2).unsqueeze(-1)) |
| if self.architecture=='gru':x,_=self.temporal(x) |
| elif self.architecture=='transformer': |
| x=x+self.position[-x.shape[1]:] |
| mask=torch.ones(x.shape[1],x.shape[1],device=x.device,dtype=torch.bool).triu(1) |
| x=self.temporal(x,mask=mask) |
| else:x=self.temporal(x) |
| current=x[:,-4:].mean(1) |
| weights=(self.key(x)*self.query(current).unsqueeze(1)).sum(-1)/(self.width**.5) |
| |
| weights=F.softmax(weights+(energy.float()+45).clamp(-20,0)/10,dim=1) |
| mean=(weights.unsqueeze(-1)*x).sum(1) |
| std=((weights.unsqueeze(-1)*(x-mean.unsqueeze(1)).square()).sum(1)+1e-5).sqrt() |
| long=self.long(pool.flatten(1)) if use_pool else torch.zeros(len(x),64,device=x.device,dtype=x.dtype) |
| z=self.head(torch.cat([current,mean,std,long],-1)) |
| direction=F.normalize(self.embedding(z).float(),dim=-1) |
| return F.normalize(self.teacher_mean+direction*self.teacher_radius,dim=-1),self.activity(z).float() |
|
|
|
|
| def batch(arrays,indices,device,context=25,pool=True): |
| x=np.asarray(arrays['q'][indices,-context:],np.float32)*np.asarray(arrays['scale'][indices,-context:],np.float32) |
| values=[x,np.asarray(arrays['pool'][indices],np.float32),np.asarray(arrays['energy'][indices,-context:],np.float32), |
| np.asarray(arrays['local_target'][indices],np.float32),np.asarray(arrays['activity'][indices],np.int64), |
| np.asarray(arrays['speaker'][indices],np.int64)] |
| return [torch.from_numpy(np.ascontiguousarray(a)).to(device) for a in values] |
|
|
|
|
| @torch.inference_mode() |
| def predict(model,arrays,indices,device,context=25,use_pool=False): |
| outputs=[];states=[] |
| for offset in range(0,len(indices),128): |
| x,p,e,_,_,_=batch(arrays,indices[offset:offset+128],device,context) |
| embedding,activity=model(x,p,e,use_pool) |
| outputs.append(embedding.cpu().numpy());states.append(activity.argmax(-1).cpu().numpy()) |
| return np.concatenate(outputs),np.concatenate(states) |
|
|
|
|
| def metrics(pred,arrays,indices,prototypes,threshold=None): |
| emb,states=pred;ids=np.asarray(arrays['speaker'][indices]);classes=np.asarray(arrays['activity'][indices]);valid=classes==0 |
| gallery=np.unique(ids[valid]);speaker=ids[valid];e=emb[valid] |
| if len(gallery)<2:raise ValueError('At least two evaluation voices required') |
| scores=e@prototypes[gallery].T;own=speaker[:,None]==gallery[None,:] |
| result=roc_metrics(scores[own],scores[~own],threshold) |
| correct=gallery[scores.argmax(1)]==speaker |
| result.update(speaker_accuracy=float(correct.mean()),macro_speaker_accuracy=float(np.mean([correct[speaker==s].mean() for s in gallery])), |
| activity_accuracy=float((states==classes).mean()),speakers=int(len(gallery)),windows=len(indices),speech_windows=int(valid.sum())) |
| target=np.asarray(arrays['local_target'][indices[valid]],np.float32) |
| target/=np.maximum(np.linalg.norm(target,axis=-1,keepdims=True),1e-8) |
| result['source_utterance_teacher_cosine']=float((e*target).sum(-1).mean()) |
| result['per_speaker']={str(int(s)):{'windows':int((speaker==s).sum()),'accuracy':float(correct[speaker==s].mean()), |
| 'cosine':float(scores[own][speaker==s].mean())} for s in gallery} |
| per_speaker=np.array([r['accuracy'] for r in result['per_speaker'].values()]) |
| bootstrap=np.random.default_rng(77).choice(per_speaker,(2000,len(per_speaker))).mean(1) |
| result['macro_accuracy_speaker_bootstrap_95ci']=[float(v) for v in np.quantile(bootstrap,[.025,.975])] |
| for i,name in enumerate(('speech','silence','overlap')): |
| use=classes==i;result[name+'_recall']=float((states[use]==i).mean()) if use.any() else None |
| result['scenarios']={} |
| scenario=np.asarray(arrays['scenario'][indices])[valid] |
| for i,name in enumerate(('clean','noise','reverb','telephone','echo')): |
| use=scenario==i |
| if use.any():result['scenarios'][name]={'windows':int(use.sum()),'accuracy':float(correct[use].mean()),'cosine':float(scores[own][use].mean())} |
| if 'emotion' in arrays: |
| from .overnight import EMOTIONS |
| emotion=np.asarray(arrays['emotion'][indices])[valid];result['emotions']={} |
| for i,name in enumerate(EMOTIONS): |
| use=emotion==i |
| if use.any():result['emotions'][name]={'windows':int(use.sum()),'accuracy':float(correct[use].mean()),'cosine':float(scores[own][use].mean())} |
| if threshold is not None: |
| enrolled=gallery[::2];known=np.isin(speaker,enrolled);score=(e@prototypes[enrolled].T).max(1) |
| result['known_acceptance']=float((score[known]>=threshold).mean()) |
| result['unknown_rejection']=float((score[~known]<threshold).mean()) |
| result['unknown_protocol']='Alternating evaluation speaker IDs enrolled; remaining identities unknown; validation-calibrated threshold' |
| return result |
|
|
|
|
| def write_report(out,results,finished=False): |
| table=''.join('<tr><td>'+html.escape(r['name'])+'</td><td>'+str(r['parameters'])+'</td><td>'+str(r['best_epoch'])+ |
| '</td><td>'+f"{r['validation']['positive_cosine']:.3f}"+'</td><td>'+f"{100*r['validation']['macro_speaker_accuracy']:.1f}%"+ |
| '</td><td>'+f"{100*r['validation']['eer']:.1f}%"+'</td></tr>' for r in results) |
| payload=json.dumps({'finished':finished,'results':results},indent=2,allow_nan=False) |
| (out/'results.json').write_text(payload) |
| (out/'report.html').write_text('''<!doctype html><meta charset="utf-8"><meta name="viewport" content="width=device-width"><title>TYGER V3 speaker study</title> |
| <style>body{font:17px system-ui;max-width:1100px;margin:40px auto;padding:20px;background:#faf8f3;color:#302f2b}td,th{padding:12px;border-bottom:1px solid #ccc;text-align:left}pre{white-space:pre-wrap;font-size:13px}a{color:#735b38}</style> |
| <h1>Speaker prediction from frozen Nemotron hidden states</h1><p>'''+('Completed evaluation.' if finished else 'Training in progress; provisional validation results.')+'''</p> |
| <table><tr><th>Candidate</th><th>Parameters</th><th>Selected epoch</th><th>Enrollment cosine</th><th>Macro speaker accuracy</th><th>EER</th></tr>'''+table+'''</table> |
| <p>480 training / 10 validation / 10 test synthetic voice identities, disjoint. Each prediction is aligned to a new 320-ms step (four 80-ms L15 states). State-pooling spans: 320, 640, 1040, and 2000 ms; optional pooled eight-second representation is reported separately. These are NOT audio-history limits: each hidden state already has the past context of its causal eight-second encoder window. No future audio. These are intermediate states from the current window, not just states cached when earlier windows ended.</p> |
| <p>Targets: frozen ECAPA 192D embeddings of clean source utterances; separate enrollment-only utterances define identity prototypes. Training combines raw and centered cosine, normalized MSE, prototype contrastive loss and an activity head. Silence/ambiguous overlap have no regression target. This is a synthetic-voice study, not secure identity verification.</p> |
| <p>Four architectures and context tests are declared before test evaluation. Selection uses validation only; held-out test is evaluated once after selection and never used to retune. Windows share utterances and voices, so they are not independent statistical trials. Random source excerpts, five augmentation scenarios and small same-split cohorts do not replace real-room or real-human evaluation. The echo scenario is a synthetic delayed source track, not a full independent robot/user interaction.</p> |
| <h2>Reproducibility and complete ablations</h2><pre>'''+html.escape(payload)+'''</pre><p><a href="/speaker-training/">Live training status</a> · <a href="/">Robby</a></p>''') |
|
|
|
|
| def run_four(work,update): |
| work=Path(work);out=work/'models';out.mkdir(exist_ok=True) |
| arrays={p.stem:np.load(p,mmap_mode='r') for p in (work/'arrays').glob('*.npy')} |
| splits={s:np.asarray(v,np.int64) for s,v in json.loads((work/'splits.json').read_text()).items()} |
| if not arrays['filled'].all():raise ValueError('Incomplete feature cache; refusing incomplete training') |
| rows=json.loads((work/'sources_manifest.json').read_text());prototypes=np.load(work/'prototypes.npy').astype(np.float32) |
| prototypes/=np.maximum(np.linalg.norm(prototypes,axis=-1,keepdims=True),1e-8) |
| trainvoices=np.array(sorted(r['speaker_idx'] for r in rows if r['split']=='train')) |
| device='cuda:0' if torch.cuda.is_available() else 'cpu';torch.set_num_threads(4) |
| if device!='cpu':torch.cuda.set_per_process_memory_fraction(.10,0) |
| train_indices=splits['train'];val_indices=splits['validation'];test_indices=splits['test'] |
| counts=np.bincount(arrays['activity'][train_indices],minlength=3).clip(1) |
| majority_activity=int(counts.argmax()) |
| class_weights=torch.tensor(counts**-.5,dtype=torch.float32,device=device);class_weights/=class_weights.mean() |
| proto=torch.from_numpy(prototypes).to(device);mean=proto[trainvoices].mean(0);radius=(proto[trainvoices]-mean).norm(dim=-1).mean() |
| contrast=F.normalize(proto[trainvoices]-mean,dim=-1) |
| lookup=torch.full((500,),-1,dtype=torch.long,device=device);lookup[torch.tensor(trainvoices,device=device)]=torch.arange(len(trainvoices),device=device) |
| results=[];rng=np.random.default_rng(1137) |
| for number,cfg in enumerate(CONFIGS): |
| seed=1730+number;torch.manual_seed(seed);random.seed(seed) |
| model=Student(cfg['architecture'],cfg['width']).to(device);model.teacher_mean.copy_(mean);model.teacher_radius.copy_(radius) |
| parameters=sum(p.numel() for p in model.parameters()) |
| if not 100000<=parameters<=1000000:raise ValueError('Student parameter budget violated: '+str(parameters)) |
| optimizer=torch.optim.AdamW(model.parameters(),lr=4e-4,weight_decay=.003) |
| latest=out/(cfg['name']+'-resume.pt');bestfile=out/(cfg['name']+'-best.pt') |
| start_epoch=0;best=-1e9;history=[];best_epoch=0 |
| if latest.exists(): |
| saved=torch.load(latest,map_location=device,weights_only=False);model.load_state_dict(saved['model']);optimizer.load_state_dict(saved['optimizer']) |
| start_epoch=saved['epoch'];history=saved['history'];best=saved['best'];best_epoch=saved['best_epoch'] |
| if 'numpy_rng' in saved:rng.bit_generator.state=saved['numpy_rng'] |
| if 'torch_rng' in saved:torch.set_rng_state(saved['torch_rng'].cpu()) |
| if device!='cpu' and 'cuda_rng' in saved:torch.cuda.set_rng_state(saved['cuda_rng'].cpu(),0) |
| start=time.perf_counter();batch_size=128 |
| for epoch in range(start_epoch,18): |
| model.train();order=rng.permutation(train_indices);running=0.;epoch_start=time.perf_counter() |
| for step,offset in enumerate(range(0,len(order),batch_size)): |
| from .overnight import guard |
| if step%100==0:guard() |
| selected=order[offset:offset+batch_size];context=int(rng.choice([4,8,13,25]));use_pool=bool(rng.random()<.25) |
| x,p,e,target,activity,speaker=batch(arrays,selected,device,context) |
| optimizer.zero_grad(set_to_none=True) |
| with torch.autocast(device_type='cuda',dtype=torch.bfloat16,enabled=device!='cpu'): |
| embedding,logits=model(x,p,e,use_pool);valid=activity==0 |
| regression=embedding.sum()*0 |
| if valid.any(): |
| a=embedding[valid];b=F.normalize(target[valid],dim=-1);ac=F.normalize(a-mean,dim=-1);bc=F.normalize(b-mean,dim=-1) |
| raw=(1-(a*b).sum(-1)).mean();center=(1-(ac*bc).sum(-1)).mean() |
| identity=F.cross_entropy(ac@contrast.T/.12,lookup[speaker[valid]]) |
| regression=.4*raw+.6*center+.05*F.mse_loss(a,b)*192+.2*identity |
| loss=regression+.2*F.cross_entropy(logits,activity,weight=class_weights) |
| if not torch.isfinite(loss):raise RuntimeError('Nonfinite loss; checkpoint retained') |
| loss.backward();nn.utils.clip_grad_norm_(model.parameters(),2);optimizer.step();running+=float(loss.detach())*len(selected) |
| if step%100==0:update(phase='Training four causal students',candidate=cfg['name'],candidate_number=number+1,epoch=epoch+1, |
| batch=step,train_seconds=time.perf_counter()-start,parameters=parameters) |
| model.eval();validation=metrics(predict(model,arrays,val_indices,device),arrays,val_indices,prototypes) |
| objective=validation['macro_speaker_accuracy']-.5*validation['eer']+.1*validation['source_utterance_teacher_cosine'] |
| history.append({'epoch':epoch+1,'loss':running/len(order),'seconds':time.perf_counter()-epoch_start,'validation':validation}) |
| if objective>best: |
| best=objective;best_epoch=epoch+1 |
| checkpoint({'model':model.state_dict(),'config':cfg,'parameters':parameters,'epoch':best_epoch,'seed':seed, |
| 'validation':validation,'format_version':3,'teacher_revision':'0f99f2d0ebe89ac095bcc5903c4dd8f72b367286'},bestfile) |
| checkpoint({'model':model.state_dict(),'optimizer':optimizer.state_dict(),'config':cfg,'epoch':epoch+1,'history':history, |
| 'best':best,'best_epoch':best_epoch,'numpy_rng':rng.bit_generator.state,'torch_rng':torch.get_rng_state(), |
| **({'cuda_rng':torch.cuda.get_rng_state(0)} if device!='cpu' else {})},latest) |
| update(phase='Training four causal students',candidate=cfg['name'],epoch=epoch+1,validation=validation) |
| if epoch>=8 and epoch+1-best_epoch>=5:break |
| saved=torch.load(bestfile,map_location=device,weights_only=False);model.load_state_dict(saved['model']);model.eval() |
| validation=saved['validation'];threshold=validation['far05_threshold'];contexts={} |
| for context in (4,8,13,25): |
| contexts[str(context*80)+'ms']=metrics(predict(model,arrays,val_indices,device,context),arrays,val_indices,prototypes,threshold) |
| contexts['2000ms_plus_8s_pool']=metrics(predict(model,arrays,val_indices,device,25,True),arrays,val_indices,prototypes,threshold) |
| |
| x,p,e,*_=batch(arrays,val_indices[:1],device,25) |
| with torch.inference_mode(): |
| for _ in range(10):model(x,p,e,False) |
| if device!='cpu':torch.cuda.synchronize() |
| t=time.perf_counter() |
| for _ in range(100):model(x,p,e,False) |
| if device!='cpu':torch.cuda.synchronize() |
| result={**cfg,'parameters':parameters,'best_epoch':saved['epoch'],'validation':validation,'validation_contexts':contexts, |
| 'student_batch1_ms':(time.perf_counter()-t)*10,'device':device,'train_seconds_this_run':time.perf_counter()-start,'history':history} |
| results.append(result);write_report(out,results) |
| del model,optimizer |
| if device!='cpu':torch.cuda.empty_cache() |
| winner=max(results,key=lambda r:r['validation']['macro_speaker_accuracy']-.5*r['validation']['eer']+.1*r['validation']['source_utterance_teacher_cosine']) |
| |
| saved=torch.load(out/(winner['name']+'-best.pt'),map_location=device,weights_only=False) |
| model=Student(saved['config']['architecture'],saved['config']['width']).to(device);model.load_state_dict(saved['model']);model.eval() |
| winner['selected']=True;threshold=winner['validation']['far05_threshold'];winner['test_contexts']={} |
| for split,ids in [('validation',val_indices),('test',test_indices)]: |
| baseline=np.broadcast_to(F.normalize(mean,dim=-1).cpu().numpy(),(len(ids),192)) |
| majority=np.full(len(ids),majority_activity) |
| source=np.asarray(arrays['local_target'][ids],np.float32) |
| source/=np.maximum(np.linalg.norm(source,axis=-1,keepdims=True),1e-8) |
| winner[split+'_controls']={ |
| 'constant_training_teacher_mean':metrics((baseline,majority),arrays,ids,prototypes,threshold), |
| 'clean_full_source_utterance_teacher':metrics((source,majority),arrays,ids,prototypes,threshold), |
| 'note':'Teacher uses clean full source audio, not equal-context/noise; activity is a training-majority constant in both controls.'} |
| for context in (4,8,13,25): |
| winner['test_contexts'][str(context*80)+'ms']=metrics(predict(model,arrays,test_indices,device,context),arrays,test_indices,prototypes,threshold) |
| winner['test_contexts']['2000ms_plus_8s_pool']=metrics(predict(model,arrays,test_indices,device,25,True),arrays,test_indices,prototypes,threshold) |
| write_report(out,results,True);update(phase='Four students complete; held-out evaluation ready',selected=winner['name'],results_path=str(out/'report.html')) |
|
|
|
|
| if __name__=='__main__': |
| from .overnight import WORK,update |
| run_four(WORK,update) |
|
|
