假设输入文本为:"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 的推理时间