File size: 10,658 Bytes
fe31b61 | 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 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """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()
|