为什么队列需要缓慢变化的键编码器？ · 实现代码

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