RuntimeError: CUDA out of memory

我的代码简化后分两个类，TokenBatchEncoder ，这个类中的self.model，就是pre_trained unixcoder , 用来做代码 embedding。其中embedding方法，根据不同的source_type做不同的embedding, 即，两种方式，一种是对每一个tokens sequence 取其vector表示。另一种是取tokens sequence 中的每一个tokens的vector表示。


class TokenBatchEncoder(nn.Module):
   def __init__(self,model,config, device=torch.device('cpu')):
        super(TokenBatchEncoder, self).__init__()
        self.model=model
        self.device=device
        self.config=config
   def embedding(self,source_ids,source_type):
       if source_type=="C":
         mask = source_ids.ne(self.config.pad_token_id)
         try:
          token_embeddings = self.model(source_ids,attention_mask = mask.unsqueeze(1) * mask.unsqueeze(2))[0]
         #sentence_embeddings = token_embeddings * mask.unsqueeze(-1)
          sentence_embeddings = (token_embeddings * mask.unsqueeze(-1)).sum(1) / mask.sum(-1).unsqueeze(-1)
          return sentence_embeddings,None
         except:
           print(torch.cuda.memory_summary())
       else:
         mask = source_ids.ne(self.config.pad_token_id)
         token_embeddings = self.model(source_ids,attention_mask = mask.unsqueeze(1) * mask.unsqueeze(2))[0]
         sentence_embeddings = token_embeddings * mask.unsqueeze(-1)
         return sentence_embeddings,mask.int()
   def tokensBatch(self,tokens,length):
       maxLen=max(length)
       batch_size=len(tokens)
       rslt=[]
       indx=[]
       for i in range(maxLen):
           index=[]
           temp=[]
           for k in range(batch_size):
               if length[k]>0:
                  code=tokens[k]
                  temp.append(code[i])
                  index.append(k)
                  length[k]=length[k]-1
           rslt.append(torch.Tensor(temp).int())
           indx.append(index)
       return rslt,indx
   def restore(self,batched,indx):
     final_rslt={}
     for i in range(len(batched)):
      index=indx[i]
      elem=batched[i]
      for m,l in zip(index,elem):
          if m not in final_rslt:
            final_rslt[m]=[l]
          else:
            final_rslt[m].append(l)
     return list(torch.stack(i) for i in final_rslt.values())
   def forward(self, tokens,batch_size,emb_type):

        embeddings=[]
        if emb_type=="C":     ＃＃＃＃ 一个batch 其中每一个是变长的tokens sequence, 
          length=[len(i) for i in tokens]
          maxlen=max(length)
          batched,index =self.tokensBatch(tokens,length)      ##  作定制的batch
          for k in range(len(batched)):
            emb,mask=self.embedding(batched[k].to(self.device),emb_type) ## embedding
            embeddings.append(emb)
          batched_emb=self.restore(embeddings,index)     ##   恢复其原来的batch
          batched_emb=torch.nn.utils.rnn.pad_sequence(batched_emb, batch_first=True)
          emb_mask=None
        else:
             batched_emb,emb_mask= self.embedding(tokens.to(self.device),emb_type)  
        return batched_emb,emb_mask
 class ClassifierModel(nn.Module):
    def __init__(self,model,config, hidden_dim,device=torch.device('cpu')):
        super(ClassifierModel, self).__init__()
        self.encoder=TokenBatchEncoder(model,config,device)
    def forward(self, view1, view2):
        """
        view1 : [[[],[]],[[],[],[]]]  变长的token sequences

        view2 :  定长的 tokens sequences.
        """
        view1_encoding,v1_mask=self.encoder(view1,len(view1),"C") 

        view2_encoding,v2_mask=self.encoder(view2,len(view2),"S")

        return  0

以上的代码，当运行几个batch后，就raise runtime error, cuda out of memory.

请教一下，有什么思路解决这个问题。查了一些网上相关的处理，包括，del variable, torch.cuda.empty_cache()等，但是没有用。

microsoft / CodeBERT

RuntimeError: CUDA out of memory #259