心法利器[107] onnx和tensorRT的bert加速方案记录
心法利器
本栏目主要和大家一起讨论近期自己学习的心得和体会，与大家一起成长。具体介绍：仓颉专项：飞机大炮我都会，利器心法我还有。
2023年新一版的文章合集已经发布，获取方式看这里：又添十万字-CS的陋室2023年文章合集来袭，更有历史文章合集，欢迎下载。
往期回顾
心法利器[102] | 大模型落地应用架构的一种模式
心法利器[103] | 大模型bad case修复方案思考
心法利器[104] | 基础RAG-向量检索模块（含代码）
心法利器[105] 基础RAG-大模型和中控模块代码（含代码）
心法利器[106] 基础RAG-调优方案
假期想着补点遗漏的知识，所以选择了模型加速这块的工作，这块我的深度肯定是不够的，不过尝试动手做做实践收获还是不小，而且形成一些自己需要的组件还是挺有用的，所以记录一下。
目录：
加速整体思路。
环境和加速前准备。
onnx加速。
tensorRT加速。
速度测试。
这里的很多代码，都有参考这篇文章，我特别摆在这里：https://blog.csdn.net/m0_37576959/article/details/127123186，感谢社区大佬的贡献。
加速整体思路
现在其实已经有大量的工具可以用来进行加速，早在bert是主流的时代，就已经有研究很多有关的技术了。今天所介绍的onnx和tensorRT也都是这个时代的产物，后续成为这个时代的主流和代表性方案，而且随着逐步迭代，他们的封装也逐步变得简单，让我们使用的难度也变低了不少。
目前这两者的加速思路其实也比较类似，即主要分为两块：
模型的重新编译，使之转化为更适用于推理的格式，并将其进行保存。
加载保存的模型，并用其进行推理。
因此，我们核心需要做的，就是上面两步的开发。
环境和加速前准备
无论是onnx和tensorrt，因为加速依赖底层硬件和操作系统，所以环境配置成了繁杂但不可绕开的工作，此处我先把我目前的环境列举出来，给大家提供参考：
windows11，16G内存，i9-13900HX，NVIDIA GeForce RTX 4070 Laptop GPU（显存8G专用+8G共享）。
CUDA Version: 12.3，CUDNN：cudnn-windows-x86_64-8.9.6.50_cuda12，python：3.9.13。
torch==2.1.2+cu121，transformers==4.33.2
onnx==1.15.0，onnxruntime==1.16.3，onnxruntime-gpu==1.17.0，基本直接pip install即可。
tensorRT：tensorrt-8.6.1.6.windows10.x86_64.cuda-12.0，tensorrt==8.6.1。
tensorrt环境配置，可以参考这篇文章：https://blog.csdn.net/KRISNAT/article/details/130789078，核心流程如下：
从NVIDIA官网上，找到自己合适的版本解压。
配置好环境变量，同时有些文件要复制到cuda内。
找到合适的python whl包，进行安装。
环境配置好了，肯定就还需要有原料了，即需要加速的模型和原始推理方案，这里我选择的是心法利器[104] | 基础RAG-向量检索模块（含代码）中提到的simcse模型。原始的加载和推理是这样的：
import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger
from tqdm import tqdm
from transformers import BertConfig, BertModel, BertTokenizer
class SimcseModel(nn.Module):
# https://blog.csdn.net/qq_44193969/article/details/126981581
def __init__(self, pretrained_bert_path, pooling="cls") -> None:
super(SimcseModel, self).__init__()
self.pretrained_bert_path = pretrained_bert_path
self.config = BertConfig.from_pretrained(self.pretrained_bert_path)
self.model = BertModel.from_pretrained(self.pretrained_bert_path, config=self.config)
self.model.eval()
# self.model = None
self.pooling = pooling
def forward(self, input_ids, attention_mask, token_type_ids):
out = self.model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
return out.last_hidden_state[:, 0]
class VectorizeModel:
def __init__(self, ptm_model_path, device = "cpu") -> None:
self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
self.model = SimcseModel(pretrained_bert_path=ptm_model_path, pooling="cls")
self.model.eval()
# self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu")
self.DEVICE = device
logger.info(device)
self.model.to(self.DEVICE)
self.pdist = nn.PairwiseDistance(2)
def predict_vec(self,query):
q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
with torch.no_grad():
q_id_input_ids = q_id["input_ids"].squeeze(1).to(self.DEVICE)
q_id_attention_mask = q_id["attention_mask"].squeeze(1).to(self.DEVICE)
q_id_token_type_ids = q_id["token_type_ids"].squeeze(1).to(self.DEVICE)
q_id_pred = self.model(q_id_input_ids, q_id_attention_mask, q_id_token_type_ids)
return q_id_pred
def predict_vec_request(self, query):
q_id_pred = self.predict_vec(query)
return q_id_pred.cpu().numpy().tolist()
def predict_sim(self, q1, q2):
q1_v = self.predict_vec(q1)
q2_v = self.predict_vec(q2)
sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
return sim.cpu().numpy().tolist()这里是有两个类，分别是SimcseModel和VectorizeModel，前者是模型类，后者是应用类，为什么这么分，在这篇文章里有提及，这里不赘述：心法利器[104] | 基础RAG-向量检索模块（含代码），当然后面用加速模型推理的时候，大家也会发现这个代码设计的优点。
另外值得强调的是，早期版本的onnx对if的算子支持的不是很好，所以大家尽量不要在模型类内增加这个，if这个还是比较常见的，所以特别说明。
onnx加速
加速的第一步，是生成新编译好的模型文件，这部分还是偏简单的，直接照着写基本就可以了。
import torch
from transformers import BertTokenizer
from src.models.vec_model.simcse_model import SimcseModel
# Reference: https://blog.csdn.net/m0_37576959/article/details/127123186
# ------------模型编译----------
# 1. 必要配置
MODEL_PATH = "C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext"
MODEL_ONNX_PATH = "./data/model_simcse_roberta_output_20240211.onnx"
DEVICE = torch.device('cuda' if torch.cuda.is_available() else "cpu")
# 2. 模型加载
tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
model = SimcseModel(pretrained_bert_path=MODEL_PATH, pooling="cls")
# OPERATOR_EXPORT_TYPE = torch._C._onnx.OperatorExportTypes.ONNX
model.eval()
model.to(DEVICE)
# 3. 格式定义
query = "你好"
encodings = tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
input_info = (encodings["input_ids"].to(DEVICE),encodings["attention_mask"].to(DEVICE),encodings["token_type_ids"].to(DEVICE))
# model(input_info)
# 4. 模型导出
output = torch.onnx.export(model,
input_info,
MODEL_ONNX_PATH,
verbose=False,
export_params=True,
# operator_export_type=OPERATOR_EXPORT_TYPE,
opset_version=12,
input_names=['input_ids', 'attention_mask', 'token_type_ids'], # 需要注意顺序！不可随意改变, 否则结果与预期不符
output_names=['output'], # 需要注意顺序, 否则在推理阶段可能用错output_names
do_constant_folding=True,
dynamic_axes={"input_ids": {0: "batch_size", 1: "length"},
"token_type_ids": {0: "batch_size", 1: "length"},
"attention_mask": {0: "batch_size", 1: "length"},
"output": {0: "batch_size"}}
)
print("Export of {} complete!".format(MODEL_ONNX_PATH))
# ------------模型校验----------
import onnxruntime as ort
import onnx
onnx_model = onnx.load(MODEL_ONNX_PATH)
onnx.checker.check_model(onnx_model)
# ------------模型校验----------
sess = ort.InferenceSession(MODEL_ONNX_PATH, providers=['CUDAExecutionProvider'])
query = "你好"
encodings = tokenizer(query, max_length = 512, truncation=True, padding="max_length", return_tensors='pt')
def to_numpy(tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
input = {
sess.get_inputs()[0].name: to_numpy(encodings["input_ids"]),
sess.get_inputs()[1].name: to_numpy(encodings["attention_mask"]),
sess.get_inputs()[2].name: to_numpy(encodings["token_type_ids"]),
}
print(sess.run(None, input_feed=input))
print(model(encodings["input_ids"].to(DEVICE),encodings["attention_mask"].to(DEVICE),encodings["token_type_ids"].to(DEVICE)))这里的必要流程我都有些写注释，应该基本足够了解了，不过还是把重点讲讲吧。
首先是转化这块，其所有流程的中心，就落在torch.onnx.export这一个函数身上，这点已经非常方便了，前面所有的工作都是为了这一个函数所需要的原料在准备，比较核心的参数只要是这几个：
model：核心需要转化的模型，torch.nn.Module肯定是可以支持的，也是比较常见的。
args：这里输入的内容，是预期模型的参数格式，描述格式本身还比较困难，格式比较多样，然而这里支持的是可以往里面塞例子，这里的第三步，也就是格式定义，就是准备了一个例子。
f：这里是指输出的路径，最终加速后的模型的路径，当然，也可以用文件对象，就是open读取的那个对象。
input_names：定义好具体模型的输入，类bert模型本身是有3个输入，名称要定义清楚，和args中一致，且注意要按照顺序。
output_names：输出的名字，这个自己定义好就行。
dynamic_axes：这里是指支持动态的变量，这里是可以指明的，动态会为速度带来一定的影响，但是也会带来较高的灵活性，一般动态的比较多的就是batch_size和句子长度了，注意这里输入和输出都得写在里面。
然后就是推理，这里的难度其实不大，基本上按照脚本走就没有什么大问题，值得注意的细节就一个，onnx模型的推理需要的输入是numpy格式，此时转化后记得要转化回来。
另外，这里有一个函数onnx.checker.check_model，这个函数是负责检验模型生成是否规范，onnx底层是用protobuf定义的，内部的各个节点的映射关系之类的都写在这里面，为了保证整体符合规范，所以出了这个函数，具体细节可以参考这篇文章：https://blog.csdn.net/qq_43456016/article/details/130256097。虽然我们这种转化比较简单，但个人还是建议在脚本里都加上。
在完成新的模型文件生成后，就可以开始推理了，直接看新的推理程序吧。
class VectorizeModel_onnx(VectorizeModel):
def __init__(self, ptm_model_path, onnx_path) -> None:
self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
self.model = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider'])
self.pdist = nn.PairwiseDistance(2)
def _to_numpy(self, tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
def predict_vec(self,query):
q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
input_feed = {
self.model.get_inputs()[0].name: self._to_numpy(q_id["input_ids"]),
self.model.get_inputs()[1].name: self._to_numpy(q_id["attention_mask"]),
self.model.get_inputs()[2].name: self._to_numpy(q_id["token_type_ids"]),
}
return torch.tensor(self.model.run(None, input_feed=input_feed)[0])
def predict_sim(self, q1, q2):
q1_v = self.predict_vec(q1)
q2_v = self.predict_vec(q2)
sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
return sim.numpy().tolist()这里的程序，为了简单，我是直接继承了上面提及的VectorizeModel，有些必要的额函数就不用重新写了。这里的加载和推理其实都参考了前面的“模型校验”中的内容了，加载用的是ort.InferenceSession，至于推理，与之不同的是，推理需要对tokenizer后的变量转为numpy的格式，另外输出这里，为了和前面的函数对齐，做好无缝切换，所以把输出的结果转化为torch.tensor了。
tensorRT加速
tensorRT的加速需要基于onnx，是需要对onnx进行进一步编译完成，流程上核心坑主要有两个：
环境配置，必须满足python版本、cuda版本等信息，而且nvidia下载速度较慢。
数据类型等细节的对齐，否则很容易失败。
详细的操作和尝试，大家可以参考这篇文章，可以说非常详细，对于详细版，大家可以看这个：https://blog.csdn.net/m0_37576959/article/details/127123186
而我想在这里聊的，是onnx本身所具有的编译tensorRT的功能，就在这行代码里：
model = ort.InferenceSession(onnx_path, providers=['CUDAExecutionProvider'])这里的providers中提供了3种：['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']，顾名思义，分别对应tensorRT、GPU、CPU三种模式，而这里的TensorrtExecutionProvider就是tensorRT编译的结果了。（https://zhuanlan.zhihu.com/p/457484536）
因此，上面对VectorizeModel_onnx就能优化为这个形式了（改个名字V2吧）：
class VectorizeModel_v2(VectorizeModel):
def __init__(self, ptm_model_path, model_path, providers=['CUDAExecutionProvider']) -> None:
# ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
self.tokenizer = BertTokenizer.from_pretrained(ptm_model_path)
self.model = ort.InferenceSession(model_path, providers=providers)
self.pdist = nn.PairwiseDistance(2)
def _to_numpy(self, tensor):
return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
def predict_vec(self,query):
q_id = self.tokenizer(query, max_length = 200, truncation=True, padding="max_length", return_tensors='pt')
input_feed = {
self.model.get_inputs()[0].name: self._to_numpy(q_id["input_ids"]),
self.model.get_inputs()[1].name: self._to_numpy(q_id["attention_mask"]),
self.model.get_inputs()[2].name: self._to_numpy(q_id["token_type_ids"]),
}
return torch.tensor(self.model.run(None, input_feed=input_feed)[0])
def predict_sim(self, q1, q2):
q1_v = self.predict_vec(q1)
q2_v = self.predict_vec(q2)
sim = F.cosine_similarity(q1_v[0], q2_v[0], dim=-1)
return sim.numpy().tolist()速度测试
有了进行加速这事，那就来对比一下加速的优化效率吧。
先来看看我的脚本：
import time,random
from tqdm import tqdm
# 加载
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
vec_model = VectorizeModel('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext', device=device)
vec_model = VectorizeModel_v2('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext',
"./data/model_simcse_roberta_output_20240211.onnx",providers=['CUDAExecutionProvider'])
vec_model = VectorizeModel_v2('C:/work/tool/huggingface/models/simcse-chinese-roberta-wwm-ext',
"./data/model_simcse_roberta_output_20240211.onnx",providers=['TensorrtExecutionProvider'])
batch_sizes = [1,2,4,8,16,32]
# 单测
# q = ["你好啊"]
# print(vec_model.predict_vec(q))
# print(vec_model.predict_sim("你好呀","你好啊"))
# 开始批跑
batch_sizes = [1,2,4,8,16]
tmp_queries = ["你好啊", "今天天气怎么样", "我要暴富"]
for b in batch_sizes:
for i in tqdm(range(100),desc="warmup"):
tmp_q = []
for i in range(b):
tmp_q.append(random.choice(tmp_queries))
vec_model.predict_vec(tmp_q)
for i in tqdm(range(1000),desc="batch_size={}".format(b)):
tmp_q = []
for i in range(b):
tmp_q.append(random.choice(tmp_queries))
vec_model.predict_vec(tmp_q)这里是4个部分，分别是加载、单测（测单独一个case）、预热和开始批跑，时间的测试用的tqdm最终的平均时间即可。事不宜迟直接给出结果吧（单位：item/s，item是指每次推理，而非每条数据）：
| batch_size | pytorch | onnx | tensorRT |
|---|---|---|---|
| 1 | 107.57 | 167.21 | 204.40 |
| 2 | 63.99 | 103.82 | 126.92 |
| 4 | 40.54 | 57.81 | 70.51 |
| 8 | 21.69 | 29.43 | 36.05 |
| 16 | 10.51 | 14.46 | 17.24 |
这里可以看到，onnx和tensorRT相比原始的pytorch模型的提升还是非常大的。
补充一个实验后的发现，tensorRT的推理中，当输入进去的数据的batch_size变化后，都会有个不短的预热时间，而在batch_size固定的那段时间，速度还是比稳定的，不知道是不是有什么bug还是这个编译情况就是如此，有了解的大佬可以在评论区里说下看有没有什么解决方案。
![心法利器[107] onnx和tensorRT的bert加速方案记录_第1张图片](http://img.e-com-net.com/image/info8/1cae63d367e54774896b264ff51a6d71.jpg)