减少 BERT 的推理时间

Posted 2023-03-29

【中文标题】减少 BERT 的推理时间

【英文标题】：Reduce inference time for BERT

【发布时间】：2021-11-10 05:44:00

【问题描述】：

我想进一步改进 BERT 的推理时间。 下面是代码：

for sentence in list(data_dict.values()):
    tokens = 'input_ids': [], 'attention_mask': []
    new_tokens = tokenizer.encode_plus(sentence, max_length=512,
                                        truncation=True, padding='max_length',
                                        return_tensors='pt',
                                        return_attention_mask=True)
    tokens['input_ids'].append(new_tokens['input_ids'][0])
    tokens['attention_mask'].append(new_tokens['attention_mask'][0])

    # reformat list of tensors into single tensor
    tokens['input_ids'] = torch.stack(tokens['input_ids'])
    tokens['attention_mask'] = torch.stack(tokens['attention_mask'])

    outputs = model(**tokens)
    embeddings = outputs[0]

有没有办法提供批次（例如在训练中）而不是整个数据集？

【问题讨论】：

如何在训练迭代中准备批次？此外，您是在示例中存储激活（缺少上下文），还是使用with torch.no_grad()mode？批处理的推理方式与训练方式相同。

感谢您的回复！对于培训，我使用的是 Huggingface 的 TrainingArguments 和 Trainer。至于激活，我不存储它们，所以 torch.no_grad()。

您是否考虑过量化您的模型以使用具有低精度数据类型的权重。您可以使用精度影响最小的低精度数据类型。

【参考方案1】：

我们可以在此处进行多项优化，这些优化（大部分）由 Huggingface 分词器原生支持。

TL;DR，优化版本将是这个版本，我在下面解释了每个更改背后的想法。

def chunker(seq, batch_size=16):
    return (seq[pos:pos + batch_size] for pos in range(0, len(seq), batch_size))

for sentence_batch in chunker(list(data_dict.values())):
    tokenized_sentences = tokenizer(sentence_batch, max_length=512,
                                truncation=True, padding=True,
                                return_tensors="pt", return_attention_mask=True)
    with torch.no_grad():
        outputs = model(**tokenized_sentences)

第一个优化是同时批量处理多个样本。为此，仔细查看标记器的实际__call__ 函数会很有帮助，请参阅here（我的粗体突出显示）：

text (str, List[str], List[List[str]]) – 要编码的序列或序列批次 [...]。

这意味着如果我们可以简单地同时传递几个样本就足够了，并且我们已经得到了容易处理的批次。我个人想指出，理论上可以一次通过整个样本列表，但也有一些我们稍后会讨论的缺点。

要实际将相当数量的样本传递给标记器，我们需要一个函数，该函数可以在一次迭代中从字典（我们的待批）中聚合多个样本。我为此使用了另一个 *** 答案，请参阅this post 了解几个有效答案。 我选择了投票率最高的答案，但请注意，这会创建和显式复制，因此可能不是最节省内存的解决方案。然后您可以简单地迭代批次，如下所示：

def chunker(seq, batch_size=16):
    return (seq[pos:pos + batch_size] for pos in range(0, len(seq), batch_size))

for sentence_batch in chunker(list(data_dict.values())):
    ...

下一个优化是您可以调用标记器的方式。您的代码通过多个步骤来完成此操作，这些步骤可以聚合到一个调用中。为了清楚起见，我还指出了您的调用中不需要哪些参数（这通常会提高您的代码可读性）。

tokenized_sentences = tokenizer(sentence_batch, max_length=512,
                                truncation=True, padding=True,
                                return_tensors="pt", return_attention_mask=True)
with torch.no_grad():  # Just to be sure
    outputs = model(**tokenized_sentences)

我还想评论一些参数的使用：

max_length=512：仅当您的值与模型的默认值max_length 不同时才需要。对于大多数模型，这将默认为 512。 return_attention_mask：也将默认为特定于模型的值，并且在大多数情况下不需要显式设置。 padding=True：如果您注意到了，这与您的版本不同，并且可以说是什么给了您最“开箱即用”的加速。通过使用padding=max_length，

每个序列将计算很多不必要的标记

，因为每个输入有 512 个标记长。对于我见过的大多数真实数据，输入往往要短得多，因此您只需要考虑批次中最长的序列长度。 padding=True 正是这样做的。对于实际的（CPU 推理）加速，我自己尝试了一些不同的序列长度，请参阅my repository on Github。值得注意的是，对于相同的 CPU 和不同的批量大小，可以实现 10 倍的加速。

编辑：我也在这里添加了torch.no_grad()，以防其他人想使用这个sn-p。我通常建议在实际受其影响的代码之前使用它，以免意外忽略任何内容。

此外，还有一些可能的优化需要您对数据样本有更深入的了解： 如果样本长度的差异非常大，如果按长度对样本进行排序（理想情况下，标记化长度，但字符长度/字数也会给你一个大概的想法），你可以获得更高的加速。这样，当将多个样本批处理在一起时，您可以最大限度地减少所需的填充量。

【参考方案2】：

也许您可能对在 CPU 上执行推理的 Intel OpenVINO 后端感兴趣？目前正在分支 https://github.com/huggingface/transformers/pull/14203

上进行工作