如何改进代码，加快转化模型的词嵌入速度？

Question

我需要为一堆不同语言模型的文档计算单词嵌入。没有问题，脚本做得很好，只是我在笔记本上工作，没有GPU，每个文本需要大约1.5s来处理，这太长了（我有成千上万的文本要处理）。

下面是我用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。

我可以想到的一些方法包括智能批处理，它是由 句子变换器在这里，你基本上是将长度相近的输入排序在一起，以避免填充到512个令牌的限制。我不知道这样做能给你带来多大的速度提升，但唯一能让你在短时间内显著提高速度的方法。

否则，如果你能获得 谷歌合作如果能在合理的时间内完成处理，你也可以利用他们的GPU环境。