LAB 07 · MoCo：历史 key、队列与两种更新

import os, tempfile
LAB_WORKDIR = tempfile.mkdtemp(prefix='lecture03-')
os.chdir(LAB_WORKDIR)
print('Ephemeral experiment workspace ready')
import numpy as np
q=np.array([1.,0.]); keys=np.array([[.8,.2],[0.,1.],[-1.,0.]])
keys=keys/np.linalg.norm(keys,axis=1,keepdims=True)
for T in [.1,.5]:
 s=keys@q/T; p=np.exp(s-s.max()); p/=p.sum(); print('T',T,'p',p.round(4))
queue=['a','b','c','d']; queue=(queue+['e','f'])[-4:]
tk=np.array([0.,2.]); tq=np.array([2.,4.]); print('queue',queue,'EMA',(0.9*tk+0.1*tq))

"""CPU image contrastive experiment. No class/color label enters pretraining."""
from pathlib import Path
import copy, hashlib, json, os, platform, time
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from sklearn.neighbors import KNeighborsClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

torch.set_num_threads(2)
SEED=20260922

SOURCE_SHA256='09bbb32bb5502db3535d89efec2045ffcd85c249b49f622390f12875060845d4'
print("source",SOURCE_SHA256,"torch",torch.__version__)

def make_data(seed=SEED):
    rng=np.random.default_rng(seed)
    yy,xx=np.mgrid[:16,:16]
    images=[]; shape=[]; color=[]; split=[]
    # Stratified before view construction. One independently generated source per ID.
    for s in range(3):
        for c in range(2):
            for i in range(140):
                cx,cy=rng.uniform(6,10,2); r=rng.uniform(3.0,4.7)
                dx,dy=xx-cx,yy-cy
                if s==0: mask=dx*dx+dy*dy<=r*r
                elif s==1: mask=np.maximum(abs(dx),abs(dy))<=r
                else: mask=abs(dx)+abs(dy)<=r*1.35
                # Red and blue have equal channel sums: grayscale erases hue by design.
                rgb=np.array([.9,.15,.15] if c==0 else [.15,.15,.9])
                bg=rng.uniform(.015,.075,(16,16,1))
                im=np.where(mask[...,None],rgb*rng.uniform(.8,1),bg)
                im+=rng.normal(0,.012,im.shape)
                images.append(np.clip(im,0,1).transpose(2,0,1))
                shape.append(s); color.append(c); split.append(0 if i<100 else 1)
    x=np.asarray(images,dtype='float32'); split=np.asarray(split)
    ids=np.arange(len(x)); train=ids[split==0]; test=ids[split==1]
    assert set(train).isdisjoint(test)
    assert len({hashlib.sha256(v.tobytes()).hexdigest() for v in x})==len(x)
    return torch.from_numpy(x),np.asarray(shape),np.asarray(color),train,test

def views(x, destroy_color=False):
    # Independent transforms of each source; no labels/scene attributes are read.
    a=x.clone(); n=len(a)
    shifts=torch.randint(-1,2,(n,2))
    for i,(dy,dx) in enumerate(shifts.tolist()):
        a[i]=torch.roll(a[i],(dy,dx),(-2,-1))
    a=a*torch.empty(n,1,1,1).uniform_(.8,1.2)
    a=a+torch.randn_like(a)*.025
    # The ablation changes only the hue swap. Every run consumes this RNG draw.
    swap=torch.rand(n)<.5
    if destroy_color: a[swap]=a[swap][:,[2,1,0]]
    return a.clamp(0,1)

class ImageEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv=nn.Sequential(nn.Conv2d(3,12,3,padding=1),nn.ReLU(),
          nn.Conv2d(12,24,3,stride=2,padding=1),nn.ReLU(),
          nn.Conv2d(24,24,3,stride=2,padding=1),nn.ReLU())
        self.readout=nn.Linear(24*4*4,48)
    def forward(self,x): return self.readout(self.conv(x).flatten(1))

class ContrastiveModel(nn.Module):
    def __init__(self):
        super().__init__(); self.encoder=ImageEncoder()
        self.projector=nn.Sequential(nn.Linear(48,48),nn.ReLU(),nn.Linear(48,24))
    def forward(self,x):
        h=self.encoder(x); z=self.projector(h)
        return h,z

def info_nce(z1,z2,tau=.2):
    z=F.normalize(torch.cat([z1,z2]),dim=1); n=len(z1)
    logits=z@z.T/tau
    logits=logits.masked_fill(torch.eye(2*n,dtype=torch.bool),-1e9)
    positive=(torch.arange(2*n)+n)%(2*n)
    return F.cross_entropy(logits,positive)

def train_simclr(initial,x,train,destroy_color,epochs=40):
    # Equal initialization, minibatch ordering, optimizer, steps, and RNG streams.
    torch.manual_seed(SEED+1)
    model=copy.deepcopy(initial); opt=torch.optim.Adam(model.parameters(),lr=.002)
    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(60):
            batch=x[ids]; a,b=views(batch,destroy_color),views(batch,destroy_color)
            loss=info_nce(model(a)[1],model(b)[1])
            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,history,delta

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

result=run_contrastive('lab07-notebook-output',mode='moco')
print('FIFO/EMA warmup',result['warmup'])
print(Path('lab07-notebook-output/queue-trace.json').read_text())
from IPython.display import Image,display
display(Image(filename='lab07-notebook-output/training-and-probes.png'))