File size: 6,718 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 | """Paid recorded-feature acquisition with exact Gaussian rank-interval value."""
from __future__ import annotations
import hashlib,json,time
import torch
from . import factorial,interval_value
from .factor_study import ROOT,PROTOCOL,verify,write
class Evidence:
def __init__(self,values):self._values=values;self.visible=values.new_full(values.shape,float('nan'));self.observed=[[] for _ in values];self.query_log=[]
def reveal(self,rows,q):
for row in rows.tolist():
assert q not in self.observed[row];self.observed[row].append(q)
self.visible[rows,q]=self._values[rows,q]
self.query_log.append(dict(rows=rows.tolist(),feature=q))
@torch.no_grad()
def episode(model,values,policy,seed,budgets=(2,4,8),cost=.01):
evidence=Evidence(values);n,d=values.shape;device=values.device
stopped=torch.zeros(n,device=device,dtype=torch.bool);counts=torch.zeros(n,device=device,dtype=torch.long)
snapshots={};cache={};start=time.perf_counter();value_evaluations=0
generator=torch.Generator(device=device).manual_seed(seed)
static_means=model['mean'];static=(static_means.var(0,unbiased=False)/model['cov'].diagonal()).argsort(descending=True,stable=True)
for step in range(1,max(budgets)+1):
groups={}
for row in range(n):
if not bool(stopped[row]):groups.setdefault(tuple(sorted(evidence.observed[row])),[]).append(row)
for key,ids in groups.items():
rows=torch.tensor(ids,device=device);visible=evidence.visible[rows]
probability=factorial.posterior(model,visible,key)
current=probability.topk(min(3,probability.shape[-1]),-1).values.sum(-1)
if key not in cache:cache[key]=factorial.conditional_channels(model,key)
means,variance=cache[key];available=[q for q in range(d) if q not in key]
if not available:stopped[rows]=True;continue
if policy.startswith('value'):
expected=[]
for q in available:
expected.append(interval_value.expected_topk(probability,means[:,q],variance[q],3));value_evaluations+=len(rows)
future,which=torch.stack(expected,-1).max(-1)
chosen=rows.new_tensor(available)[which]
stop=(future-current<=cost) if policy.endswith('stop') else torch.zeros_like(current,dtype=torch.bool)
elif policy.startswith('variance'):
mu=means[:,available];average=probability@mu
score=((probability@mu.square())-average.square())/variance[available]
chosen=rows.new_tensor(available)[score.argmax(-1)];stop=torch.zeros_like(current,dtype=torch.bool)
if policy.endswith('stop'):
for q in chosen.unique().tolist():
selected=chosen==q;future=interval_value.expected_topk(probability[selected],means[:,q],variance[q],3)
stop[selected]=future-current[selected]<=cost;value_evaluations+=int(selected.sum())
elif policy=='static_forced':
q=next(q for q in static.tolist() if q in available);chosen=rows.new_full((len(rows),),q);stop=torch.zeros_like(current,dtype=torch.bool)
else:
chosen=rows.new_tensor(available)[torch.randint(len(available),(len(rows),),device=device,generator=generator)];stop=torch.zeros_like(current,dtype=torch.bool)
stopped[rows[stop]]=True
for q in chosen[~stop].unique().tolist():
selected=rows[(chosen==q)&~stop];evidence.reveal(selected,q);counts[selected]+=1
if step in budgets:
probs=values.new_zeros((n,len(model['prior'])))
groups={}
for row in range(n):groups.setdefault(tuple(sorted(evidence.observed[row])),[]).append(row)
for key,ids in groups.items():
rows=torch.tensor(ids,device=device);probs[rows]=factorial.posterior(model,evidence.visible[rows],key)
snapshots[step]=dict(probability=probs,queries=counts.clone())
torch.cuda.synchronize();return snapshots,dict(seconds=time.perf_counter()-start,queried_features=evidence.query_log,
conditional_masks_materialized=len(cache),exact_value_evaluations=value_evaluations,
future_observation_draws=0,withheld_values_never_read_by_inference=True)
def main(device='cuda:1'):
verify();start=time.perf_counter();torch.cuda.set_device(device);torch.backends.cuda.matmul.allow_tf32=False
config=json.loads(PROTOCOL.read_text());manifest=json.loads((ROOT/'prepared/manifest.json').read_text());settings=config['secondary_acquisition'];all_results=[]
for series in manifest['series']:
if not series['primary_eligible']:continue
p=torch.load(ROOT/'prepared'/series['file'],map_location=device,weights_only=True)
rows=torch.where(p['masks']['unseen'])[0].tolist()
rows=sorted(rows,key=lambda i:hashlib.sha256(('20781:'+series['records'][i]['id']).encode()).hexdigest())[:32]
selected=torch.tensor(rows,device=device);sizes=[len(v) for v in series['levels']];true=factorial.joint_index(p['h'][selected],sizes)
for name in settings['world_models']:
path=ROOT/'models'/series['name']/f'n64_s{config["sample_seeds"][0]}'/f'{name}.pt'
checkpoint=torch.load(path,map_location=device,weights_only=True);model=checkpoint['model'];x=(p['x'][selected]-checkpoint['mean'])/checkpoint['scale']
for policy in settings['policies']:
outputs,trace=episode(model,x,policy,20801,tuple(settings['budgets']),settings['per_scalar_recorded_query_cost'])
for budget,out in outputs.items():
credit=factorial.topk_credit(out['probability'],true,3)
all_results.append(dict(series=series['name'],system=series['system'],model=name,policy=policy,budget=budget,
selected_rows=selected,groups=p['groups'][selected],top3=credit,queries=out['queries'],
net_utility=credit-settings['per_scalar_recorded_query_cost']*out['queries'],trace=trace))
print(json.dumps(dict(series=series['name'],model=name,policy=policy,seconds=trace['seconds'])),flush=True)
torch.save(dict(results=all_results),ROOT/'acquisition.pt');write('acquisition_execution.json',dict(seconds=time.perf_counter()-start,
peak_gpu_bytes=torch.cuda.max_memory_allocated(device),cells=len(all_results),subset_fixed_before_predictions=True,
complete_primary_diagnostic_cohort_retained=True,no_globally_optimal_multistep_claim=True))
verify();print('FULL-HYPOTHESIS ACQUISITION COMPLETE',flush=True)
if __name__=='__main__':main()
|