| """Post-evaluation checks of frozen models; no fitting or recipe selection.""" |
| from pathlib import Path |
| import argparse,hashlib,importlib.util,itertools,json,sys,time |
| PROCESS_START=time.perf_counter() |
| import torch |
|
|
| def aggregate(value,days): |
| if days is None:return float(value.mean()) |
| return float(torch.stack([value[days==d].mean() for d in days.unique()]).mean()) |
|
|
| def macro_combination(value,codes,days): |
| if days is None:days=torch.zeros_like(codes) |
| return float(torch.stack([torch.stack([value[(days==d)&(codes==h)].mean() |
| for h in codes[days==d].unique()]).mean() for d in days.unique()]).mean()) |
|
|
| def sensor_checks(replay,device): |
| output=[];supports=[];maximum=0.;comparisons=[] |
| with torch.no_grad(): |
| for cohort,sizes in [('hydraulic',[3,4,3,4]),('temperature',[10,3])]: |
| p=replay.load(f'data/{cohort}/prepared',device) |
| if cohort=='hydraulic':p.update(p['masks']) |
| records=replay.load(f'results/{cohort}/predictions',device)['results'] |
| for r in records: |
| n,seed,name=r['samples'],r['seed'],r['model'] |
| checkpoint=replay.load(f'models/{cohort}/n{n}_s{seed}/{name}',device) |
| fit=checkpoint['fit_ids'];held=r['held_rows'];mask=r['unseen'] |
| assert bool(p['source'][fit].all()) and not bool(p['unseen'][fit].any()) |
| prob=replay.predict(checkpoint['model'],(p['x'][held]-checkpoint['mean'])/checkpoint['scale']) |
| truth=replay.factorial.joint_index(p['h'][held],sizes) |
| scores={f'top{k}':replay.factorial.topk_credit(prob,truth,k) for k in (1,3)} |
| scores.update(deterministic_top1=(prob.argmax(-1)==truth).double(), |
| deterministic_top3=(prob.topk(3,-1).indices==truth[:,None]).any(-1).double()) |
| scores.update(nll=-prob.gather(1,truth[:,None]).squeeze(-1).clamp_min(1e-300).log(), |
| confidence=prob.max(-1).values,top3_mass=prob.topk(3,-1).values.sum(-1)) |
| for key,value in scores.items(): |
| error=float((value-r['metrics'][key]).abs().max());maximum=max(maximum,error) |
| assert error<1e-8,(cohort,n,seed,name,key,error) |
| codes=truth[mask];days=r['days'][mask] if cohort=='temperature' else None |
| row=dict(cohort=cohort,samples=n,seed=seed,model=name,records=int(mask.sum()), |
| observed_combinations=len(codes.unique()), |
| primary_top1=aggregate(scores['top1'][mask],days), |
| primary_top3=aggregate(scores['top3'][mask],days), |
| deterministic_top3=aggregate(scores['deterministic_top3'][mask],days), |
| macro_combination_top3=macro_combination(scores['top3'][mask],codes,days), |
| top3_probability_mass=aggregate(scores['top3_mass'][mask],days)) |
| distance=(p['h'][held][mask,None]!=p['h'][fit][None]).sum(-1).min(-1).values |
| assert bool((distance>=1).all()) |
| row['nearest_fitting_combination']=[dict(hamming_distance=int(d),records=int((distance==d).sum()), |
| top3=aggregate(scores['top3'][mask][distance==d],days[distance==d] if days is not None else None)) |
| for d in distance.unique()] |
| top4=prob[mask].topk(4,-1).values.clamp_min(1e-300) |
| row['median_third_fourth_log_probability_margin']=float((top4[:,2].log()-top4[:,3].log()).median()) |
| row['per_combination']=[dict(code=int(h),levels=p['h'][held][mask][(codes==h).nonzero()[0,0]].tolist(), |
| records=int((codes==h).sum()),top3=aggregate(scores['top3'][mask][codes==h],days[codes==h] if days is not None else None)) |
| for h in codes.unique()] |
| output.append(row) |
| if name=='factorized_neural': |
| supports.append(dict(cohort=cohort,samples=n,seed=seed, |
| factor_level_counts=[torch.bincount(p['h'][fit,i],minlength=k).tolist() for i,k in enumerate(sizes)], |
| source_combinations=len(replay.factorial.joint_index(p['h'][fit],sizes).unique()))) |
| for n in [64,128]: |
| for control in ['factorized_classifier','additive_full','factorized_lda']: |
| a={r['seed']:r for r in output if r['cohort']==cohort and r['samples']==n and r['model']=='factorized_neural'} |
| b={r['seed']:r for r in output if r['cohort']==cohort and r['samples']==n and r['model']==control} |
| comparisons.append(dict(cohort=cohort,samples=n,control=control, |
| paired_seed_gains=[dict(seed=s,primary=a[s]['primary_top3']-b[s]['primary_top3'], |
| macro_combination=a[s]['macro_combination_top3']-b[s]['macro_combination_top3']) for s in sorted(a)])) |
| summary=[] |
| for cohort in ['hydraulic','temperature']: |
| for n in [64,128]: |
| for name in sorted({r['model'] for r in output}): |
| rows=[r for r in output if r['cohort']==cohort and r['samples']==n and r['model']==name] |
| summary.append(dict(cohort=cohort,samples=n,model=name,observed_combinations=rows[0]['observed_combinations'], |
| primary_top3=sum(r['primary_top3'] for r in rows)/len(rows), |
| deterministic_top3=sum(r['deterministic_top3'] for r in rows)/len(rows), |
| macro_combination_top3=sum(r['macro_combination_top3'] for r in rows)/len(rows), |
| seed_top3_min=min(r['primary_top3'] for r in rows),seed_top3_max=max(r['primary_top3'] for r in rows), |
| mean_top3_probability_mass=sum(r['top3_probability_mass'] for r in rows)/len(rows))) |
| assert len(output)==132 |
| return dict(checkpoints=len(output),maximum_replay_error=maximum,summary=summary, |
| seed_comparisons=comparisons,fitting_support=supports,individual_fits=output) |
|
|
| def enumerate_trees(device,budget): |
| """All depth-limited binary-query trees, independently of the DP recurrence.""" |
| outcomes=list(itertools.product([0,1],repeat=2));terminal=[[1.,-2.],[-2.,1.],[0.,0.]] |
| trees=[([a]*4,[0]*4) for a in range(3)] |
| for first in [0,1]: |
| other=1-first |
| branch_options=[('stop',a,a) for a in range(3)] |
| if budget==2:branch_options += [('query',a,b) for a in range(3) for b in range(3)] |
| for branches in itertools.product(branch_options,repeat=2): |
| actions=[];queries=[] |
| for outcome in outcomes: |
| kind,a,b=branches[outcome[first]] |
| actions.append(a if kind=='stop' or outcome[other]==0 else b) |
| queries.append(1+int(kind=='query')) |
| trees.append((actions,queries)) |
| assert len(trees)==(21 if budget==1 else 291) |
| utilities=torch.tensor([[[terminal[a][h] for a in actions] for h in range(2)] for actions,_ in trees],device=device,dtype=torch.float64) |
| queries=torch.tensor([q for _,q in trees],device=device,dtype=torch.float64)[:,None].expand(-1,2,-1) |
| return utilities.flatten(1),queries.flatten(1) |
|
|
| def planner_check(replay,device): |
| core=replay.core;worlds=replay.load('data/controlled/worlds',device);maximum=0.;comparisons=0 |
| names=['identity','query_joint','posterior_clusters','posterior_bins','future_marginal','raw_information','raw_entropy','hash_code'] |
| trees={b:enumerate_trees(device,b) for b in [1,2]} |
| with torch.no_grad(): |
| for stream in range(3): |
| fitted=core.estimate(core.count_samples(worlds['samples'][:,stream,:1024],6),'support_mixture') |
| for name in names: |
| code,k=core.encode(fitted,name);law=core.push(fitted,code,k);policy=core.plan(law) |
| flat=law.flatten(2) |
| for budget in [1,2]: |
| reward,queries=trees[budget] |
| for ci,cost in enumerate([0.,.01,.05]): |
| expected=(flat@(reward-cost*queries).T).max(-1).values.sum(-1) |
| error=float((expected-policy['expected'][:,budget,ci]).abs().max()) |
| maximum=max(maximum,error);comparisons+=len(law) |
| assert error<1e-12,(name,budget,cost,error) |
| return dict(worlds=192,streams=3,encodings=len(names),utility_comparisons=comparisons, |
| depth_one_trees=21,depth_two_trees=291,maximum_absolute_error=maximum, |
| verification='All deterministic binary-query trees; separate from the dynamic-programming recurrence.') |
|
|
| def main(): |
| parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--release',type=Path) |
| parser.add_argument('--device',default='cuda:1');parser.add_argument('--output',type=Path,required=True);a=parser.parse_args() |
| root=a.release or Path(__file__).resolve().parents[1] |
| if not (root/'reproduce.py').is_file():root=root/'release/huggingface' |
| root=root.resolve();sys.path.insert(0,str(root)) |
| spec=importlib.util.spec_from_file_location('frozen_replay',root/'reproduce.py') |
| replay=importlib.util.module_from_spec(spec);spec.loader.exec_module(replay) |
| assert torch.cuda.is_available() and a.device.startswith('cuda') |
| t=time.perf_counter();torch.cuda.set_device(a.device);torch.cuda.init();torch.cuda.synchronize() |
| init=time.perf_counter()-t;torch.backends.cuda.matmul.allow_tf32=False |
| t=time.perf_counter();sensor=sensor_checks(replay,a.device);torch.cuda.synchronize();sensor_seconds=time.perf_counter()-t |
| t=time.perf_counter();planner=planner_check(replay,a.device);torch.cuda.synchronize();planner_seconds=time.perf_counter()-t |
| model_hashes={str(p.relative_to(root)):hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted((root/'models').rglob('*')) if p.is_file()} |
| report=dict(post_evaluation_descriptive_analysis=True,new_training=False,new_thresholds=False,new_confirmation=False, |
| sensor=sensor,enumerative_planner_check=planner,model_files_sha256=model_hashes, |
| compute=dict(cuda_initialization_seconds=init,sensor_diagnostics_seconds=sensor_seconds, |
| planner_verification_seconds=planner_seconds,process_seconds_after_standard_library_import=time.perf_counter()-PROCESS_START, |
| peak_cuda_bytes=torch.cuda.max_memory_allocated(a.device),device=torch.cuda.get_device_name(a.device))) |
| a.output.parent.mkdir(parents=True,exist_ok=True);a.output.write_text(json.dumps(report,indent=2)+'\n') |
| print(json.dumps(dict(checkpoints=sensor['checkpoints'],maximum_replay_error=sensor['maximum_replay_error'], |
| planner=planner,compute=report['compute']),indent=2)) |
|
|
| if __name__=='__main__':main() |
|
|
