# f_q, f_k: encoder networks for query and key 
# queue: dictionary as a queue of K keys (CxK) 
# m: momentum 
# t: temperature  

f_k.params = f_q.params # initialize 
for x in loader: # load a minibatch x with N samples 
	x_q = aug(x) # a randomly augmented version 
	x_k = aug(x) # another randomly augmented version  
	q = f_q.forward(x_q) # queries: NxC 
	k = f_k.forward(x_k) # keys: NxC 
	k = k.detach() # no gradient to keys  
	
	# positive logits: Nx1 
	l_pos = bmm(q.view(N,1,C), k.view(N,C,1))  # 相当于把batch中每个正样本对之间求了cosine临近
	
	# negative logits: NxK 
	l_neg = mm(q.view(N,C), queue.view(C,K))  
	
	# logits: Nx(1+K) 
	logits = cat([l_pos, l_neg], dim=1)  
	
	# contrastive loss, Eqn.(1) 
		labels = zeros(N) # positives are the 0-th，将识别的类别视为0,可以直接使用CrossEntropyLoss
	loss = CrossEntropyLoss(logits/t, labels)  
	
	# SGD update: query network 
	loss.backward() 
	update(f_q.params)  
	
	# momentum update: key network 
	f_k.params = m*f_k.params+(1-m)*f_q.params  
	
	# update dictionary 
	enqueue(queue, k) # enqueue the current minibatch 
	dequeue(queue) # dequeue the earliest minibatch

# image_encoder - ResNet or Vision Transformer 
# text_encoder - CBOW or Text Transformer 
# I[n, h, w, c] - minibatch of aligned images 
# T[n, l] - minibatch of aligned texts 
# W_i[d_i, d_e] - learned proj of image to embed 
# W_t[d_t, d_e] - learned proj of text to embed 
# t - learned temperature parameter  

# extract feature representations of each modality 
I_f = image_encoder(I) #[n, d_i] 
T_f = text_encoder(T) #[n, d_t]  

# joint multimodal embedding [n, d_e] 
I_e = l2_normalize(np.dot(I_f, W_i), axis=1) T
_e = l2_normalize(np.dot(T_f, W_t), axis=1)  

# scaled pairwise cosine similarities [n, n] 
logits = np.dot(I_e, T_e.T) * np.exp(t)  

# symmetric loss function 
labels = np.arange(n) 
loss_i = cross_entropy_loss(logits, labels, axis=0) 
loss_t = cross_entropy_loss(logits, labels, axis=1) 
loss = (loss_i + loss_t)/2