如何改进代码以加速 transformer 模型的词嵌入？

Question

我需要为一堆具有不同语言模型的文档计算词嵌入。没问题，脚本运行良好，除了我在没有 GPU 的笔记本上工作，每个文本需要大约 1.5 秒才能处理，这太长了（我有数千个文本要处理）。

以下是我使用 pytorch 和 transformers lib 的方式：

import torch
from transformers import CamembertModel, CamembertTokenizer

docs = [text1, text2, ..., text20000]
tok = CamembertTokenizer.from_pretrained('camembert-base')
model = CamembertModel.from_pretrained('camembert-base', output_hidden_states=True)
# let try with a batch size of 64 documents
docids = [tok.encode(
  doc, max_length=512, return_tensors='pt', pad_to_max_length=True) for doc in docs[:64]]
ids=torch.cat(tuple(docids))
device = 'cuda' if torch.cuda.is_available() else 'cpu' # cpu in my case...
model = model.to(device)
ids = ids.to(device)
model.eval()
with torch.no_grad():
    out = model(input_ids=ids)
# 103s later...

有人对提高速度有什么想法或建议吗？

Answer 1

我不认为有一种简单的方法可以在不使用 GPU 的情况下显着提高速度。

我能想到的一些方法包括 Sentence-Transformers 使用的智能批处理，您基本上将相似长度的输入排序在一起，以避免填充到完整的 512 个标记限制。我不确定这会给你带来多少加速，但这是你可以在短时间内显着提高它的唯一方法。

否则，如果您可以访问 Google colab，如果可以在合理的时间内完成处理，您也可以利用他们的 GPU 环境。

如何改进代码以加速 transformer 模型的词嵌入？

How to improve code to speed up word embedding with transformer models?

nlp

pytorch

word-embedding

huggingface-transformers

bert-language-model