| """Matched compact direct classifiers, including source-only calibration.""" |
| from __future__ import annotations |
| import torch |
| from torch import nn |
| from .factorial import hypotheses |
|
|
|
|
| def train(x,h,sizes,seed,nonlinear,steps=400): |
| torch.cuda.manual_seed(seed);d=x.shape[1];out=sum(sizes) |
| if nonlinear: |
| width=(19984-out)//(d+1+out) |
| net=nn.Sequential(nn.Linear(d,width,device=x.device),nn.SiLU(),nn.Linear(width,out,device=x.device)) |
| else:net=nn.Linear(d,out,device=x.device) |
| opt=torch.optim.AdamW(net.parameters(),lr=.003,weight_decay=.01) |
| counts=[torch.bincount(h[:,i],minlength=k).double()+1 for i,k in enumerate(sizes)] |
| with torch.enable_grad(): |
| for _ in range(steps): |
| opt.zero_grad(set_to_none=True);logits=net(x.float());loss=x.new_zeros(());offset=0 |
| for i,size in enumerate(sizes): |
| weights=(1/counts[i]).float();weights/=weights.mean() |
| loss=loss+nn.functional.cross_entropy(logits[:,offset:offset+size],h[:,i],weight=weights)/len(sizes);offset+=size |
| loss.backward();opt.step() |
| return net |
|
|
|
|
| def fit(x,h,sizes,kind,seed,steps=400): |
| nonlinear=kind=='factorized_classifier';n=len(x);out=sum(sizes) |
| network=train(x,h,sizes,seed,nonlinear,steps) |
| oof=x.new_zeros((n,out));fold=torch.arange(n,device=x.device)%4 |
| for f in range(4): |
| partial=train(x[fold!=f],h[fold!=f],sizes,seed+101+f,nonlinear,steps) |
| with torch.no_grad():oof[fold==f]=partial(x[fold==f].float()).double() |
| temperatures=[];offset=0 |
| for f,size in enumerate(sizes): |
| options=x.new_tensor([.5,1.,2.,4.]);logits=oof[:,None,offset:offset+size]/options[None,:,None] |
| loss=-logits.log_softmax(-1).gather(-1,h[:,f,None,None].expand(-1,len(options),1)).squeeze(-1).mean(0) |
| temperatures.append(options[loss.argmin()]);offset+=size |
| parameters=sum(p.numel() for p in network.parameters());assert parameters+len(sizes)<=20000 |
| return dict(kind=kind,sizes=tuple(sizes),state_dict=network.state_dict(), |
| nonlinear=nonlinear,input_dimensions=x.shape[1],grid=hypotheses(sizes,x.device), |
| temperatures=torch.stack(temperatures),neural_parameters=parameters+len(sizes),adaptation_fits=5) |
|
|
|
|
| def posterior(model,x,observed): |
| |
| assert len(list(observed))==x.shape[1] |
| state=model['state_dict'] |
| if model['nonlinear']: |
| hidden=nn.functional.silu(x.float()@state['0.weight'].T+state['0.bias']) |
| logits=hidden@state['2.weight'].T+state['2.bias'] |
| else:logits=x.float()@state['weight'].T+state['bias'] |
| out=logits.new_zeros((len(x),len(model['grid'])));offset=0 |
| for f,size in enumerate(model['sizes']): |
| logp=(logits[:,offset:offset+size]/model['temperatures'][f]).log_softmax(-1) |
| out+=logp[:,model['grid'][:,f]];offset+=size |
| return out.double().softmax(-1) |
|
|
