ONNX Runtime 推理 BERT 模型

假设输入文本为:"My name is John"

1. 使用BertTokenizer进行tokenize和encoding:

python
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

text = "My name is John"
encoded_input = tokenizer.encode_plus(text, max_length=10, truncation=True, return_tensors='pt')

print(encoded_input)
# 输出:
{'input_ids': tensor([[101, 1045, 1005, 2310, 2028, 1012, 102, 0, 0, 0]]), 
'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]),
'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 0, 0, 0]])}

input_ids: [CLS] 101, My 1045, name 1005, is 2310, John 2028, [SEP] 1012, [PAD] 0

101 - 表示[CLS]特殊token,放在序列开头

1045 - 表示token化后词汇"My"对应的id1005 - 表示token化后词汇"name"对应的id

2310 - 表示token化后词汇"is"对应的id

2028 - 表示token化后词汇"John"对应的id

1012 - 表示[SEP]特殊token,表示句子结束

0 - 在max_length不足的情况下,使用0进行padding

所以这个input_ids其实就是对原始文本"My name is John"进行了以下处理:

1. 使用BertTokenizer进行分词token化,映射到词汇表的id

2. 在开头添加[CLS]特殊token,结尾添加[SEP]

3. 使用0 pad到最大长度10

最后形成的输入就是:[101, 1045, 1005, 2310, 2028, 1012, 0, 0, 0, 0]

这可以表示为:[CLS] 101, My 1045, name 1005, is 2310, John 2028, [SEP] 1012, [PAD] 0

所以input_ids就是文本序列经过BERT模型所需预处理后的数字化表示。

token_type_ids: 全0,表示属于同一句子

attention_mask: 对应位置的注意力mask

2. 保存这些编码后的输入,作为ONNX Runtime的输入:

python
input_ids = encoded_input['input_ids'].numpy()
token_type_ids = encoded_input['token_type_ids'].numpy() 

# input_ids: [[101, 1045, 1005, 2310, 2028, 1012, 0, 0, 0, 0]]
# token_type_ids: [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]

np.savez('input.npz', 
         input_ids=input_ids,
         token_type_ids=token_type_ids)

3. ONNX Runtime推理:

python
import onnxruntime as rt

sess = rt.InferenceSession("bert.onnx")

data = np.load('input.npz')
# 还需补充attention_mask
input_ids = data['input_ids']
token_type_ids = data['token_type_ids']

outputs = sess.run(None, {
    'input_ids': input_ids,
    'token_type_ids': token_type_ids
})

4. 处理输出logits,得到预测mask位置的词:

python 
import torch

logits = outputs[0] 

# logits shape: (1, 10, 30522)
softmax = torch.softmax(torch.tensor(logits), dim=-1)  
values, indices = torch.topk(softmax, 5)

for i in indices[0]:
    print(tokenizer.decode([i])) 
# mask位置预测: 'is' 'was' 'were' 'has' 'had'

ONNXRuntime 和 TensorRT 做推理时间比较：

具体的代码流程：

1. 测量 ONNX 模型在 ONNX Runtime 的推理时间

2. 使用 ONNX Parser 将 ONNX 模型转换为 TensorRT 引擎

    - 创建 Builder,Network,Parser 对象
    - 解析 ONNX 模型生成 TensorRT network
    - 使用 Builder 生成 TensorRT 引擎(plan文件)

3. 部署 TensorRT 引擎进行推理

    - 创建 Execution Context
    - 分配显存
    - 执行推理
    - 测量推理时间

4. 比较 ONNX Runtime 和 TensorRT 的推理时间