冻结参数以后才使用标签 · 实现代码

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 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