[LAB 07 · MoCo:历史 key、队列与两种更新](/course/notebooks/0d1b8f6fbf6a476c93ffcc169f2ec1f4.html) # 为什么队列需要缓慢变化的键编码器? · 实现代码 [独立阅读](/course/notebook-blocks/ebf4defb3d804d1084f0dd9f3a91cfe9.html) 按问题继续 真实FIFO键、EMA、梯度更新及低于随机基线的保留结果 展开完整实现代码 ``` 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()) ```