两个目标，两个分母 · 实现代码

def train_ae(x,train,kind,epochs=40):
    torch.manual_seed(SEED+(kind=='masked'))
    model=ConvAE(2 if kind=='masked' else 1);initial=copy.deepcopy(model)
    opt=torch.optim.Adam(model.parameters(),lr=.003);history=[]
    before=next(model.encoder.parameters()).detach().clone()
    for epoch in range(epochs):
        model.train();losses=[]
        for ids in torch.as_tensor(train)[torch.randperm(len(train))].split(64):
            clean=x[ids]
            if kind=='denoise':
                pred,_=model(noise_corrupt(clean));loss=F.mse_loss(pred,clean)
            else:
                mask=patch_mask(len(ids),.5)
                inp=torch.cat([clean*(1-mask),mask],dim=1)
                pred,_=model(inp);loss=(((pred-clean)**2)*mask).sum()/mask.sum()
            opt.zero_grad();loss.backward();opt.step();losses.append(float(loss.detach()))
        history.append(float(np.mean(losses)))
    delta=float((next(model.encoder.parameters()).detach()-before).norm());assert delta>0
    return model,initial,history,delta

def frozen_probe(model,x,y,train,test,masked=False):
    model.eval()
    with torch.no_grad():
        # Clean image readout; mask channel zero. No reconstruction targets at inference.
        inp=torch.cat([x,torch.zeros_like(x)],1) if masked else x
        h=model(inp)[1]
        a=F.normalize(h,dim=1).numpy()
    probe=KNeighborsClassifier(n_neighbors=5,metric='cosine').fit(a[train],y[train])
    return float(probe.score(a[test],y[test]))

def mask_index_warmup():
    rng=np.random.default_rng(7);tokens=np.arange(16);shuffle=rng.permutation(16)
    restore=np.argsort(shuffle);keep=shuffle[:4]
    shuffled=np.concatenate([tokens[keep],np.full(12,-1)])
    restored=shuffled[restore]
    assert np.array_equal(restored[keep],tokens[keep]);assert (restored==-1).sum()==12
    proxy=[]
    for visible in ([0,3,9,14],[0,5,10,15,3,6,9,12]):
        mask=np.ones(16,bool);mask[visible]=False
        pred=np.full(16,tokens[visible].mean())
        proxy.append({'visible':len(visible),'masked':int(mask.sum()),'constant_predictor_masked_mse':float(((pred[mask]-tokens[mask])**2).mean())})
    return {'shuffle':shuffle.tolist(),'ids_restore':restore.tolist(),'restored':restored.tolist(),'constant_proxy':proxy}