心法利器[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还是这个编译情况就是如此，有了解的大佬可以在评论区里说下看有没有什么解决方案。