import torch import torch.nn as nn from prettytable import PrettyTable from torch.nn.modules.activation import Tanh import copy import logging logger = logging.getLogger(__name__) from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup, RobertaConfig, RobertaModel, RobertaTokenizer) def whitening_torch_final(embeddings): mu = torch.mean(embeddings, dim=0, keepdim=True) cov = torch.mm((embeddings - mu).t(), embeddings - mu) u, s, vt = torch.svd(cov) W = torch.mm(u, torch.diag(1/torch.sqrt(s))) embeddings = torch.mm(embeddings - mu, W) return embeddings class BaseModel(nn.Module): def __init__(self, ): super().__init__() def model_parameters(self): table = PrettyTable() table.field_names = ["Layer Name", "Output Shape", "Param #"] table.align["Layer Name"] = "l" table.align["Output Shape"] = "r" table.align["Param #"] = "r" for name, parameters in self.named_parameters(): if parameters.requires_grad: table.add_row([name, str(list(parameters.shape)), parameters.numel()]) return table class Model(BaseModel): def __init__(self, encoder): super(Model, self).__init__() self.encoder = encoder def forward(self, code_inputs=None, nl_inputs=None): # code_inputs [bs, seq] if code_inputs is not None: outputs = self.encoder(code_inputs,attention_mask=code_inputs.ne(1))[0] #[bs, seq_len, dim] outputs = (outputs*code_inputs.ne(1)[:,:,None]).sum(1)/code_inputs.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， return torch.nn.functional.normalize(outputs, p=2, dim=1) else: outputs = self.encoder(nl_inputs,attention_mask=nl_inputs.ne(1))[0] outputs = (outputs*nl_inputs.ne(1)[:,:,None]).sum(1)/nl_inputs.ne(1).sum(-1)[:,None] return torch.nn.functional.normalize(outputs, p=2, dim=1) class Multi_Loss_CoCoSoDa( BaseModel): def __init__(self, base_encoder, args, mlp=False): super(Multi_Loss_CoCoSoDa, self).__init__() self.K = args.moco_k self.m = args.moco_m self.T = args.moco_t dim= args.moco_dim # create the encoders # num_classes is the output fc dimension self.code_encoder_q = base_encoder self.code_encoder_k = copy.deepcopy(base_encoder) self.nl_encoder_q = base_encoder # self.nl_encoder_q = RobertaModel.from_pretrained("roberta-base") self.nl_encoder_k = copy.deepcopy(self.nl_encoder_q) self.mlp = mlp self.time_score= args.time_score self.do_whitening = args.do_whitening self.do_ineer_loss = args.do_ineer_loss self.agg_way = args.agg_way self.args = args for param_q, param_k in zip(self.code_encoder_q.parameters(), self.code_encoder_k.parameters()): param_k.data.copy_(param_q.data) # initialize param_k.requires_grad = False # not update by gradient for param_q, param_k in zip(self.nl_encoder_q.parameters(), self.nl_encoder_k.parameters()): param_k.data.copy_(param_q.data) # initialize param_k.requires_grad = False # not update by gradient # create the code queue torch.manual_seed(3047) torch.cuda.manual_seed(3047) self.register_buffer("code_queue", torch.randn(dim,self.K )) self.code_queue = nn.functional.normalize(self.code_queue, dim=0) self.register_buffer("code_queue_ptr", torch.zeros(1, dtype=torch.long)) # create the masked code queue self.register_buffer("masked_code_queue", torch.randn(dim, self.K )) self.masked_code_queue = nn.functional.normalize(self.masked_code_queue, dim=0) self.register_buffer("masked_code_queue_ptr", torch.zeros(1, dtype=torch.long)) # create the nl queue self.register_buffer("nl_queue", torch.randn(dim, self.K )) self.nl_queue = nn.functional.normalize(self.nl_queue, dim=0) self.register_buffer("nl_queue_ptr", torch.zeros(1, dtype=torch.long)) # create the masked nl queue self.register_buffer("masked_nl_queue", torch.randn(dim, self.K )) self.masked_nl_queue= nn.functional.normalize(self.masked_nl_queue, dim=0) self.register_buffer("masked_nl_queue_ptr", torch.zeros(1, dtype=torch.long)) @torch.no_grad() def _momentum_update_key_encoder(self): """ Momentum update of the key encoder % key encoder的Momentum update """ for param_q, param_k in zip(self.code_encoder_q.parameters(), self.code_encoder_k.parameters()): param_k.data = param_k.data * self.m + param_q.data * (1. - self.m) for param_q, param_k in zip(self.nl_encoder_q.parameters(), self.nl_encoder_k.parameters()): param_k.data = param_k.data * self.m + param_q.data * (1. - self.m) if self.mlp: for param_q, param_k in zip(self.code_encoder_q_fc.parameters(), self.code_encoder_k_fc.parameters()): param_k.data = param_k.data * self.m + param_q.data * (1. - self.m) for param_q, param_k in zip(self.nl_encoder_q_fc.parameters(), self.nl_encoder_k_fc.parameters()): param_k.data = param_k.data * self.m + param_q.data * (1. - self.m) @torch.no_grad() def _dequeue_and_enqueue(self, keys, option='code'): # gather keys before updating queue # keys = concat_all_gather(keys) batch_size = keys.shape[0] if option == 'code': code_ptr = int(self.code_queue_ptr) assert self.K % batch_size == 0 # for simplicity # replace the keys at ptr (dequeue and enqueue) try: self.code_queue[:, code_ptr:code_ptr + batch_size] = keys.T except: print(code_ptr) print(batch_size) print(keys.shape) exit(111) code_ptr = (code_ptr + batch_size) % self.K # move pointer ptr->pointer self.code_queue_ptr[0] = code_ptr elif option == 'masked_code': masked_code_ptr = int(self.masked_code_queue_ptr) assert self.K % batch_size == 0 # for simplicity # replace the keys at ptr (dequeue and enqueue) try: self.masked_code_queue[:, masked_code_ptr:masked_code_ptr + batch_size] = keys.T except: print(masked_code_ptr) print(batch_size) print(keys.shape) exit(111) masked_code_ptr = (masked_code_ptr + batch_size) % self.K # move pointer ptr->pointer self.masked_code_queue_ptr[0] = masked_code_ptr elif option == 'nl': nl_ptr = int(self.nl_queue_ptr) assert self.K % batch_size == 0 # for simplicity # replace the keys at ptr (dequeue and enqueue) self.nl_queue[:, nl_ptr:nl_ptr + batch_size] = keys.T nl_ptr = (nl_ptr + batch_size) % self.K # move pointer ptr->pointer self.nl_queue_ptr[0] = nl_ptr elif option == 'masked_nl': masked_nl_ptr = int(self.masked_nl_queue_ptr) assert self.K % batch_size == 0 # for simplicity # replace the keys at ptr (dequeue and enqueue) self.masked_nl_queue[:, masked_nl_ptr:masked_nl_ptr + batch_size] = keys.T masked_nl_ptr = (masked_nl_ptr + batch_size) % self.K # move pointer ptr->pointer self.masked_nl_queue_ptr[0] = masked_nl_ptr def forward(self, source_code_q, source_code_k, nl_q,nl_k): """ Input: im_q: a batch of query images im_k: a batch of key images Output: logits, targets """ if not self.args.do_multi_lang_continue_pre_train: # logger.info(".do_multi_lang_continue_pre_train") outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[0] code_q = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， code_q = torch.nn.functional.normalize(code_q, p=2, dim=1) # compute query features for nl outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[0] # queries: NxC bs*feature_dim nl_q = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None] nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1) code2nl_logits = torch.einsum("ab,cb->ac", code_q,nl_q ) # loss = self.loss_fct(scores*20, torch.arange(code_inputs.size(0), device=scores.device)) code2nl_logits /= self.T # label code2nl_label = torch.arange(code2nl_logits.size(0), device=code2nl_logits.device) return code2nl_logits,code2nl_label, None, None if self.agg_way == "avg": # compute query features for source code outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[0] code_q = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， code_q = torch.nn.functional.normalize(code_q, p=2, dim=1) # compute query features for nl outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[0] # queries: NxC bs*feature_dim nl_q = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None] nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1) # compute key features with torch.no_grad(): # no gradient to keys self._momentum_update_key_encoder() # update the key encoder # shuffle for making use of BN # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k) # masked code outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1))[0] # keys: NxC code_k = (outputs*source_code_k.ne(1)[:,:,None]).sum(1)/source_code_k.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， code_k = torch.nn.functional.normalize( code_k, p=2, dim=1) # masked nl outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1))[0] # keys: bs*dim nl_k = (outputs*nl_k.ne(1)[:,:,None]).sum(1)/nl_k.ne(1).sum(-1)[:,None] nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1) elif self.agg_way == "cls_pooler": # logger.info(self.agg_way ) # compute query features for source code outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1))[1] code_q = torch.nn.functional.normalize(code_q, p=2, dim=1) # compute query features for nl outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1))[1] # queries: NxC bs*feature_dim nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1) # compute key features with torch.no_grad(): # no gradient to keys self._momentum_update_key_encoder() # update the key encoder # shuffle for making use of BN # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k) # masked code outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1))[1] # keys: NxC code_k = torch.nn.functional.normalize( code_k, p=2, dim=1) # masked nl outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1))[1] # keys: bs*dim nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1) elif self.agg_way == "avg_cls_pooler": # logger.info(self.agg_way ) outputs = self.code_encoder_q(source_code_q, attention_mask=source_code_q.ne(1)) code_q_cls = outputs[1] outputs = outputs[0] code_q_avg = (outputs*source_code_q.ne(1)[:,:,None]).sum(1)/source_code_q.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， code_q = code_q_cls + code_q_avg code_q = torch.nn.functional.normalize(code_q, p=2, dim=1) # compute query features for nl outputs= self.nl_encoder_q(nl_q, attention_mask=nl_q.ne(1)) nl_q_cls = outputs[1] outputs= outputs[0] # queries: NxC bs*feature_dim nl_q_avg = (outputs*nl_q.ne(1)[:,:,None]).sum(1)/nl_q.ne(1).sum(-1)[:,None] nl_q = nl_q_avg + nl_q_cls nl_q = torch.nn.functional.normalize(nl_q, p=2, dim=1) # compute key features with torch.no_grad(): # no gradient to keys self._momentum_update_key_encoder() # update the key encoder # shuffle for making use of BN # im_k, idx_unshuffle = self._batch_shuffle_ddp(im_k) # masked code outputs = self.code_encoder_k(source_code_k, attention_mask=source_code_k.ne(1)) code_k_cls = outputs[1] # keys: NxC outputs = outputs[0] code_k_avg = (outputs*source_code_k.ne(1)[:,:,None]).sum(1)/source_code_k.ne(1).sum(-1)[:,None] # None作为ndarray或tensor的索引作用是增加维度， code_k = code_k_cls + code_k_avg code_k = torch.nn.functional.normalize( code_k, p=2, dim=1) # masked nl outputs = self.nl_encoder_k(nl_k, attention_mask=nl_k.ne(1)) nl_k_cls = outputs[1] # keys: bs*dim outputs = outputs[0] nl_k_avg = (outputs*nl_k.ne(1)[:,:,None]).sum(1)/nl_k.ne(1).sum(-1)[:,None] nl_k = nl_k_cls + nl_k_avg nl_k = torch.nn.functional.normalize(nl_k, p=2, dim=1) # ## do_whitening # if self.do_whitening: # code_q = whitening_torch_final(code_q) # code_k = whitening_torch_final(code_k) # nl_q = whitening_torch_final(nl_q) # nl_k = whitening_torch_final(nl_k) ## code vs nl code2nl_pos = torch.einsum('nc,bc->nb', [code_q, nl_q]) # negative logits: NxK code2nl_neg = torch.einsum('nc,ck->nk', [code_q, self.nl_queue.clone().detach()]) # logits: Nx(n+K) code2nl_logits = torch.cat([self.time_score*code2nl_pos, code2nl_neg], dim=1) # apply temperature code2nl_logits /= self.T # label code2nl_label = torch.arange(code2nl_logits.size(0), device=code2nl_logits.device) ## code vs masked nl code2maskednl_pos = torch.einsum('nc,bc->nb', [code_q, nl_k]) # negative logits: NxK code2maskednl_neg = torch.einsum('nc,ck->nk', [code_q, self.masked_nl_queue.clone().detach()]) # logits: Nx(n+K) code2maskednl_logits = torch.cat([self.time_score*code2maskednl_pos, code2maskednl_neg], dim=1) # apply temperature code2maskednl_logits /= self.T # label code2maskednl_label = torch.arange(code2maskednl_logits.size(0), device=code2maskednl_logits.device) ## nl vs code # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1) nl2code_pos = torch.einsum('nc,bc->nb', [nl_q, code_q]) # negative logits: bsxK nl2code_neg = torch.einsum('nc,ck->nk', [nl_q, self.code_queue.clone().detach()]) # nl2code_logits: bsx(n+K) nl2code_logits = torch.cat([self.time_score*nl2code_pos, nl2code_neg], dim=1) # apply temperature nl2code_logits /= self.T # label nl2code_label = torch.arange(nl2code_logits.size(0), device=nl2code_logits.device) ## nl vs masked code # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1) nl2maskedcode_pos = torch.einsum('nc,bc->nb', [nl_q, code_k]) # negative logits: bsxK nl2maskedcode_neg = torch.einsum('nc,ck->nk', [nl_q, self.masked_code_queue.clone().detach()]) # nl2code_logits: bsx(n+K) nl2maskedcode_logits = torch.cat([self.time_score*nl2maskedcode_pos, nl2maskedcode_neg], dim=1) # apply temperature nl2maskedcode_logits /= self.T # label nl2maskedcode_label = torch.arange(nl2maskedcode_logits.size(0), device=nl2maskedcode_logits.device) #logit 4*bsx(1+K) inter_logits = torch.cat((code2nl_logits, code2maskednl_logits, nl2code_logits ,nl2maskedcode_logits ), dim=0) # labels: positive key indicators # inter_labels = torch.zeros(inter_logits.shape[0], dtype=torch.long).cuda() inter_labels = torch.cat((code2nl_label, code2maskednl_label, nl2code_label, nl2maskedcode_label), dim=0) if self.do_ineer_loss: # logger.info("do_ineer_loss") ## code vs masked code code2maskedcode_pos = torch.einsum('nc,bc->nb', [code_q, code_k]) # negative logits: NxK code2maskedcode_neg = torch.einsum('nc,ck->nk', [code_q, self.masked_code_queue.clone().detach()]) # logits: Nx(n+K) code2maskedcode_logits = torch.cat([self.time_score*code2maskedcode_pos, code2maskedcode_neg], dim=1) # apply temperature code2maskedcode_logits /= self.T # label code2maskedcode_label = torch.arange(code2maskedcode_logits.size(0), device=code2maskedcode_logits.device) ## nl vs masked nl # nl2code_pos = torch.einsum('nc,nc->n', [nl_q, code_k]).unsqueeze(-1) nl2maskednl_pos = torch.einsum('nc,bc->nb', [nl_q, nl_k]) # negative logits: bsxK nl2maskednl_neg = torch.einsum('nc,ck->nk', [nl_q, self.masked_nl_queue.clone().detach()]) # nl2code_logits: bsx(n+K) nl2maskednl_logits = torch.cat([self.time_score*nl2maskednl_pos, nl2maskednl_neg], dim=1) # apply temperature nl2maskednl_logits /= self.T # label nl2maskednl_label = torch.arange(nl2maskednl_logits.size(0), device=nl2maskednl_logits.device) #logit 6*bsx(1+K) inter_logits = torch.cat((inter_logits, code2maskedcode_logits, nl2maskednl_logits), dim=0) # labels: positive key indicators # inter_labels = torch.zeros(inter_logits.shape[0], dtype=torch.long).cuda() inter_labels = torch.cat(( inter_labels, code2maskedcode_label, nl2maskednl_label ), dim=0) # dequeue and enqueue self._dequeue_and_enqueue(code_q, option='code') self._dequeue_and_enqueue(nl_q, option='nl') self._dequeue_and_enqueue(code_k, option='masked_code') self._dequeue_and_enqueue(nl_k, option='masked_nl') return inter_logits, inter_labels, code_q, nl_q