如何从CodeBERT/GraphCodeBERT等预训练LLM前4层获取Embeddings以提速？

核心修改思路 #

• 开启模型的output_hidden_states=True参数，获取所有层的隐藏状态输出
• 提取前4层Transformer的输出（注：模型返回的hidden_states元组中，索引0是词嵌入层输出，索引1到N对应第1到第N层Transformer的输出）
• 对指定层的隐藏状态做均值池化，生成最终Embedding

修改后的代码 #

def gen_embeddings(self, code):
    tokenized_input_pos = self.tokenizer(code, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        # 启用hidden_states输出，获取所有层的隐藏状态
        output = self.model(**tokenized_input_pos, output_hidden_states=True)
    # 取第4层Transformer的输出（索引4对应前4层计算后的结果：索引0是embedding，1-4是前4层）
    fourth_layer_hidden = output.hidden_states[4]
    # 对token维度做均值池化
    embedding = fourth_layer_hidden.mean(dim=1).squeeze().tolist()
    # 保持原返回逻辑：单样本返回列表，多样本直接返回
    if len(code) == 1:
        return [embedding]
    else:
        return embedding

额外提速建议 #

• 如果模型结构支持，可以直接裁剪模型只保留前4层，彻底避免后续层的计算开销：

  # 以BERT类模型为例，裁剪前4层Transformer编码器
  self.model.encoder.layer = self.model.encoder.layer[:4]
  # 裁剪后正常前向，直接取last_hidden_state即可
  

• 确保torch.no_grad()上下文管理器正确使用，杜绝不必要的梯度计算
• 批量处理时尽量保持batch_size稳定，减少tokenizer的动态padding开销

内容的提问来源于stack exchange，提问作者not-a-bot