Python NLP中Bert问答模块处理速度过慢的原因排查及优化咨询

你遇到的15-20分钟超长耗时，绝大多数是函数实现逻辑不合理导致，硬件性能只是次要影响因素，按以下方案优化后可以把耗时降到秒级。

核心问题点 #

当前代码最大的不合理设计是每次调用函数都重新加载模型和分词器：
BertForQuestionAnswering.from_pretrained和BertTokenizer.from_pretrained放在了find_answer函数内部，每次调用都会重新读取、加载超过1GB的预训练模型权重，仅这一步在普通CPU上就要耗时10分钟以上，属于完全可以避免的重复开销。

可落地的优化项 #

• 模型和分词器改为全局预加载，只在程序启动时加载一次
• 启用GPU加速，bert大模型在CPU上推理本身就非常慢，只要有NVIDIA显卡就能获得几十倍的速度提升
• 增加输入长度截断，bert原生最大支持512个token的输入，过长的输入不仅无用还会大幅增加计算量
• 推理阶段关闭梯度计算，减少内存占用同时提升推理速度

优化后的代码示例 #

import torch
from transformers import BertForQuestionAnswering, BertTokenizer

# 全局只加载一次模型和分词器
device = "cuda" if torch.cuda.is_available() else "cpu"
model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad').to(device)
tokenizer = BertTokenizer.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')

@torch.no_grad() # 关闭梯度计算
def find_answer(question,answer_text):
    # 编码时增加截断和最大长度限制
    input_ids = tokenizer.encode(question, answer_text, truncation=True, max_length=512)
    tokens = tokenizer.convert_ids_to_tokens(input_ids)
    sep_index = input_ids.index(tokenizer.sep_token_id)
    num_seg_a = sep_index + 1
    num_seg_b = len(input_ids) - num_seg_a
    segment_ids = [0] * num_seg_a + [1] * num_seg_b
    assert len(segment_ids) == len(input_ids)
    # 数据移到对应设备
    outputs = model(
        torch.tensor([input_ids]).to(device),
        token_type_ids=torch.tensor([segment_ids]).to(device),
        return_dict=True
    )
    start_scores = outputs.start_logits
    end_scores = outputs.end_logits
    answer_start = torch.argmax(start_scores)
    answer_end = torch.argmax(end_scores)
    answer = ' '.join(tokens[answer_start:answer_end + 1])
    return answer

硬件优化建议 #

如果没有NVIDIA显卡，只能用CPU运行的话，更换核心数更多、主频更高的CPU可以获得一定速度提升，但提升幅度远低于使用GPU。优化后的代码在普通消费级GTX1660显卡上，单样本推理耗时可以控制在1秒以内。

内容的提问来源于stack exchange，提问作者Qureshi