冻结评价回答另一个问题 · 实现代码

def evaluate(model,x,ys,yc,train,test):
    model.eval()
    with torch.no_grad(): h,z=model(x)
    scores={}
    for space,features in [('h',h),('z',z)]:
        a=F.normalize(features,dim=1).numpy()
        for target,y in [('shape',ys),('color',yc)]:
            knn=KNeighborsClassifier(n_neighbors=5,metric='cosine')
            knn.fit(a[train],y[train])
            # Standardization and classifier fit use TRAIN only; test used once for reporting.
            linear=make_pipeline(StandardScaler(),LogisticRegression(C=1,max_iter=1000))
            linear.fit(a[train],y[train])
            scores[f'{space}_{target}_knn']=float(knn.score(a[test],y[test]))
            scores[f'{space}_{target}_linear']=float(linear.score(a[test],y[test]))
    with torch.no_grad():
        # Counterfactual hue swap only: same geometry and pixels, channels exchanged.
        ah,az=model(x[test][:,[2,1,0]])
        for space,original,other in [('h',h[test],ah),('z',z[test],az)]:
            scores[f'{space}_hue_swap_cosine']=float(F.cosine_similarity(original,other).mean())
    return scores

def moco_warmup():
    q=np.array([1.,0.]); keys=np.array([[.8,.2],[0.,1.],[-1.,0.]])
    keys/=np.linalg.norm(keys,axis=1,keepdims=True)
    probs={}
    for T in [.1,.5]:
        s=keys@q/T; p=np.exp(s-s.max());p/=p.sum();probs[str(T)]=p.tolist()
    queue=['a','b','c','d']; queue=(queue+['e','f'])[-4:]
    ema=.9*np.array([0.,2.])+.1*np.array([2.,4.])
    assert queue==['c','d','e','f']; assert np.allclose(ema,[.2,2.2])
    return {'temperature_probabilities':probs,'FIFO':queue,'EMA':ema.tolist()}

def train_moco(initial,x,train,epochs=15):
    # Small genuine image implementation: query grads, detached EMA keys, FIFO negatives.
    torch.manual_seed(SEED+2)
    query=copy.deepcopy(initial); key=copy.deepcopy(initial)
    for p in key.parameters(): p.requires_grad=False
    opt=torch.optim.Adam(query.parameters(),lr=.002)
    queue=torch.empty((0,24)); queue_ids=torch.empty(0,dtype=torch.long)
    history=[]; trace=[]; m=.95; capacity=240
    before=next(query.encoder.parameters()).detach().clone()
    for epoch in range(epochs):
        losses=[]
        for ids in torch.as_tensor(train)[torch.randperm(len(train))].split(60):
            a,b=views(x[ids]),views(x[ids])
            q=F.normalize(query(a)[1],dim=1)
            with torch.no_grad():
                for pk,pq in zip(key.parameters(),query.parameters()):pk.mul_(m).add_(pq,alpha=1-m)
                k=F.normalize(key(b)[1],dim=1)
            # Warmup uses detached current keys. Later batches use a real FIFO dictionary.
            if not len(queue):
                logits=q@k.T/.2; labels=torch.arange(len(ids))
                loss=F.cross_entropy(logits,labels)
            else:
                positive=(q*k).sum(1,keepdim=True)
                negative=q@queue.T
                # A source already in queue is NOT treated as its own negative.
                negative=negative.masked_fill(ids[:,None]==queue_ids[None,:],-1e9)
                logits=torch.cat([positive,negative],dim=1)/.2
                loss=F.cross_entropy(logits,torch.zeros(len(ids),dtype=torch.long))
            opt.zero_grad();loss.backward();opt.step()
            with torch.no_grad():
                queue=torch.cat([queue,k])[-capacity:];queue_ids=torch.cat([queue_ids,ids])[-capacity:]
            losses.append(float(loss.detach()))
            if len(trace)<5:trace.append({'step':len(trace),'queue_size':len(queue),'first_source_id':int(queue_ids[0]),'last_source_id':int(queue_ids[-1])})
        history.append(float(np.mean(losses)))
    assert all(p.grad is None for p in key.parameters())
    return query,key,queue,queue_ids,history,trace,float((next(query.encoder.parameters()).detach()-before).norm())

def plot_results(out,x,initial,models,results,histories):
    fig,ax=plt.subplots(3,6,figsize=(12,6))
    ids=[0,100,140,280,420,700]
    torch.manual_seed(SEED+6)
    for j,i in enumerate(ids):
        ax[0,j].imshow(x[i].permute(1,2,0));ax[0,j].set_title(f'source {i}')
        paired_state=torch.get_rng_state()
        for row,erase in [(1,False),(2,True)]:
            torch.set_rng_state(paired_state)
            ax[row,j].imshow(views(x[i:i+1],erase)[0].permute(1,2,0))
    for row,label in enumerate(['source','view: hue kept','view: hue can swap']):ax[row,0].set_ylabel(label)
    for a in ax.flat:a.set_xticks([]);a.set_yticks([])
    fig.tight_layout();fig.savefig(out/'views.png',dpi=140);plt.close(fig)
    fig,axes=plt.subplots(1,3,figsize=(14,4))
    for name,hist in histories.items():axes[0].plot(np.arange(1,len(hist)+1),hist,label=name)
    axes[0].set(xlabel='epoch',ylabel='InfoNCE (batch mean)',title='Training objective');axes[0].legend(fontsize=8)
    names=list(results)
    for ax,task in zip(axes[1:],['shape','color']):
        pos=np.arange(len(names)); width=.34
        for off,space in [(-width/2,'h'),(width/2,'z')]:
            ax.bar(pos+off,[results[n][f'{space}_{task}_knn'] for n in names],width,label=space)
        ax.set_xticks(pos,names,rotation=20);ax.set_ylim(0,1.05);ax.set_title(f'Frozen cosine 5-NN: {task}')
        ax.set_ylabel('held-out accuracy');ax.legend()
    fig.tight_layout();fig.savefig(out/'training-and-probes.png',dpi=140);plt.close(fig)

def run_contrastive(out_dir='lab06-output',mode='color'):
    started=time.time();out=Path(out_dir);out.mkdir(parents=True,exist_ok=True)
    torch.manual_seed(SEED)
    x,ys,yc,train,test=make_data()
    np.savez_compressed(out/'inputs.npz',images=x.numpy(),shape=ys,color=yc,train_ids=train,test_ids=test)
    initial=ContrastiveModel(); torch.save(initial.state_dict(),out/'random-init.pt')
    results={'random':evaluate(initial,x,ys,yc,train,test)};histories={};models={};deltas={}
    if mode=='color':
        for name,erase in [('keep-hue',False),('swap-hue',True)]:
            model,hist,delta=train_simclr(initial,x,train,erase)
            results[name]=evaluate(model,x,ys,yc,train,test);histories[name]=hist;models[name]=model;deltas[name]=delta
            torch.save(model.state_dict(),out/f'{name}.pt')
    else:
        model,hist,delta=train_simclr(initial,x,train,False)
        results['simclr']=evaluate(model,x,ys,yc,train,test);histories['simclr']=hist;models['simclr']=model;deltas['simclr']=delta
        torch.save(model.state_dict(),out/'simclr.pt')
        query,key,queue,queue_ids,hist,trace,delta=train_moco(initial,x,train)
        results['moco-toy']=evaluate(query,x,ys,yc,train,test);histories['moco-toy']=hist;models['moco-toy']=query;deltas['moco-toy']=delta
        torch.save({'query':query.state_dict(),'key':key.state_dict(),'queue':queue,'queue_ids':queue_ids},out/'moco-toy.pt')
        (out/'queue-trace.json').write_text(json.dumps(trace,indent=2))
    config={'seed':SEED,'mode':mode,'image_shape':[3,16,16],'train_n':len(train),'test_n':len(test),'batch_size':60,'epochs_simclr':40,'epochs_moco':15,'tau':.2,'adam_lr':.002,'h_dim':48,'z_dim':24,'labels_used_only_for_probes':True,'knn':{'k':5,'metric':'cosine'},'linear_probe':{'regularization_C':1,'standardizer_fit':'train-only'},'test_tuning':'none; fixed protocol before first run'}
    (out/'config.json').write_text(json.dumps(config,indent=2))
    plot_results(out,x,initial,models,results,histories)
    payload={'code_sha256':globals().get('SOURCE_SHA256') or hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),'kind':'local CPU full training + frozen held-out probes','utc':time.strftime('%Y-%m-%dT%H:%M:%SZ',time.gmtime()),'seconds':time.time()-started,'torch':torch.__version__,'numpy':np.__version__,'python':platform.python_version(),'scores':results,'loss_history':histories,'encoder_weight_l2_change':deltas,'warmup':moco_warmup(),'inputs_sha256':hashlib.sha256((out/'inputs.npz').read_bytes()).hexdigest()}
    (out/'results.json').write_text(json.dumps(payload,indent=2))
    print(json.dumps({'scores':results,'encoder_weight_l2_change':deltas,'seconds':payload['seconds']},indent=2))
    return payload