如何更快地运行 bert 的预训练模型词嵌入?

数据挖掘 nlp 熊猫 词嵌入 伯特
2021-09-24 02:07:50

我正在尝试使用 microsoft/pubmedbert 获取临床数据的词嵌入。我有 360 万行文本。将文本转换为 10k 行的向量大约需要 30 分钟。因此,对于 360 万行,大约需要 - 180 小时(大约 8 天)。

有什么方法可以加快这个过程吗?

我的代码 -

from transformers import AutoTokenizer
from transformers import pipeline
model_name = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
tokenizer = AutoTokenizer.from_pretrained(model_name)
classifier = pipeline('feature-extraction',model=model_name, tokenizer=tokenizer)

def lambda_func(row):
    tokens = tokenizer(row['notetext'])
    if len(tokens['input_ids'])>512:
        tokens = re.split(r'\b', row['notetext'])
        tokens= [t for t in tokens if len(t) > 0 ]
        row['notetext'] = ''.join(tokens[:512])
    row['vectors'] = classifier(row['notetext'])[0][0]        
    return row

def process(progress_notes):     
    progress_notes = progress_notes.apply(lambda_func, axis=1)
    return progress_notes

progress_notes = process(progress_notes)
vectors_2d = np.reshape(progress_notes['vectors'].to_list(), (vectors_length, vectors_breadth))
vectors_df = pd.DataFrame(vectors_2d)

我的 progress_notes 数据框看起来像 -

progress_notes = pd.DataFrame({'id':[1,2,3],'progressnotetype':['Nursing Note', 'Nursing Note', 'Administration Note'], 'notetext': ['Patient\'s skin is grossly intact with exception of skin tear to r inner elbow and r lateral lower leg','Patient with history of Afib with RVR. Patient is incontinent of bowel and bladder.','Give 2 tablet by mouth every 4 hours as needed for Mild to moderate Pain Not to exceed 3 grams in 24 hours']})

注意 - 1)我在 aws ec2 实例 r5.8x large(32 个 CPU)上运行代码 - 我尝试使用多处理,但代码进入死锁,因为 bert 占用了我所有的 cpu 内核。

2个回答

我认为主要问题是您如何使用 BERT,因为您正在逐句处理文本。相反,您应该以小批量的方式将输入提供给模型:

用于 NLP 的神经网络不仅要一次接收一个句子,还要接收多个句子。句子以单个整数张量(令牌 ID)堆叠在一起,维度为句子数×序列长度。由于句子有不同的长度,通常批次的序列长度是其中最长序列的长度,所有比它短的句子都用填充标记填充。您可以在这里查看huggingface 库中的批处理

使用面向批处理的处理可以让您从批处理中所有句子的并行处理中获益。

问题在于,虽然 HuggingFace Transformer 支持面向批处理的训练,但在某些情况下它不支持面向批处理的推理。例如FeatureExtractionPipeline,它可以根据需要提取令牌嵌入,但不支持批处理(与TableQuestionAnsweringPipeline具有sequential参数的不同)。

这样,为了进行面向批处理的推理,您需要手动输入数据,而不是依赖管道 API。您可以在此线程中找到有关如何执行此操作的示例

如果您尝试使用 GPU 而不是 CPU,使用面向批处理的处理也是实现性能提升的关键。

如果您决定继续使用 CPU,请确保您的 Pytorch 构建使用 MKL-DNN,这是 CPU 的主要性能提升器。您可以查看此线程以了解如何操作。如果您不使用它,请安装包含它的较新版本。

有几件事可以使您的代码更快:

  • 停止使用熊猫。Pandas 不是为大规模、快速的文本处理而设计的。最好切换到Apache Arrow之类的东西,它专为现代硬件上的高效分析操作而设计。

  • 重构代码以避免转换为内存效率低的数据类型。progress_notes['vectors'].to_list()转换为将使用大量内存的 Python 列表。

  • 重构代码以避免在数据类型之间来回转换。替换列表推导并''.join使用已编译的regex编译的正则表达式处理将比等效的 Python 代码快得多。

  • 重构代码以用内联代码替换函数。Python 为每个函数调用创建一个新堆栈。内联代码将更节省内存。