赞
踩
基于TensorRT的BERT实时自然语言理解(下)
BERT Inference with TensorRT
请参阅Python脚本bert_inference.py还有详细的Jupyter notebook BERT_TRT.ipynb在sample文件夹中进行推理过程的逐步描述和演练。在本节中,让我们回顾几个关键参数和概念,以便使用TensorRT进行推理。
BERT(更具体地说是编码器层)使用以下参数来控制其操作:
Batch size
Sequence Length
Number of attention heads
这些参数的值取决于所选的BERT模型,用于设置TensorRT计划文件(执行引擎)的配置参数。
对于每个编码器,还要指定隐藏层的数量和注意头的大小。您也可以从Tensorflow检查点文件中读取上述所有参数。
由于我们所使用的BERT模型已经针对SQuAD数据集上的问答下游任务进行了微调,网络的输出(即输出完全连接层)将是一段文本,答案出现在文章中(在示例中称为h_output)。一旦我们生成了TensorRT引擎,我们就可以对其进行序列化,然后在TensorRT运行时中使用它。
在推理过程中,我们异步地执行从CPU到GPU的内存复制,反之亦然,从而分别将张量导入和移出GPU内存。异步内存复制操作通过与设备和主机之间的内存复制操作重叠计算来隐藏内存传输的延迟。异步内存复制和内核执行如图3所示。
Figure 3: TensorRT Runtime process
The inputs to the BERT model, which are shown in Figure 3, include:
input_ids: tensor with token ids of paragraph concatenated along with question that is used as input for inference
segment_ids: distinguishes between passage and question
input_mask: indicates which elements in the sequence are tokens, and which ones are padding elements
输出(start_logits)和end_logits结束逻辑表示答案的范围,网络根据问题在文章中预测答案的范围。
在Jupyter笔记本中,我们使用了作为TensorRT插件实现的层的加速
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。