心法利器[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, BertTokenizerclass SimcseModel(nn.Module):# https://blog.csdn.net/qq_44193969/article/details/126981581def __init__(self, pretrained_bert_path, pooling="cls") -> None:super(SimcseModel, self).__init__()self.pretrained_bert_path = pretrained_bert_pathself.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 = Noneself.pooling = poolingdef 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 = devicelogger.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_preddef 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()

这里是有两个类,分别是SimcseModelVectorizeModel,前者是模型类,后者是应用类,为什么这么分,在这篇文章里有提及,这里不赘述:心法利器[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_namesdo_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 onnxonnx_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_sizepytorchonnxtensorRT
1107.57167.21204.40
263.99103.82126.92
440.5457.8170.51
821.6929.4336.05
1610.5114.4617.24

这里可以看到,onnx和tensorRT相比原始的pytorch模型的提升还是非常大的。

补充一个实验后的发现,tensorRT的推理中,当输入进去的数据的batch_size变化后,都会有个不短的预热时间,而在batch_size固定的那段时间,速度还是比稳定的,不知道是不是有什么bug还是这个编译情况就是如此,有了解的大佬可以在评论区里说下看有没有什么解决方案。

10e4893054f80e1aa02cb2febb3b0469.png

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.hqwc.cn/news/469636.html

如若内容造成侵权/违法违规/事实不符,请联系编程知识网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

山西电力市场日前价格预测【2024-02-10】

日前价格预测 预测说明: 如上图所示,预测明日(2024-02-10)山西电力市场全天平均日前电价为126.73元/MWh。其中,最高日前电价为302.95元/MWh,预计出现在08:15。最低日前电价为0.00元/MWh,预计出…

jmeter遇到连接数据库的问题

jmeter连接mysql或者oracle简单,但是连接过inceptor吗? 上货 1、下载驱动inceptor 5.1.2.jar包 2、在添加驱动那里导入 3、在JBC request中的写法 PS:没什么可说的

ZISUOJ 2022年算法基础公选课练习四(Map)

说明: 博主为了提早预习数据结构和C的一些知识,自己琢磨外加查阅资料所写的代码,题目来源于22年初的学院老师组织的算法基础公选课的练习。我的代码甚至思路肯定存在许多不足和错误,欢迎大家批评指正。 题目列表: 问题…

【C++】类和对象(四)

前言:在类和对象中,我们走过了十分漫长的道路,今天我们将进一步学习类和对象,类和对象这块荆棘地很长,各位一起加油呀。 💖 博主CSDN主页:卫卫卫的个人主页 💞 👉 专栏分类:高质量&a…

旋转位置编码原理及代码

旋转位置编码原理及代码 旋转位置编码 RoPE(Rotary Positional Encoding)当位置发生偏移时,只需要旋转角度外推性,指大模型输入长度超过预训练文本长度时,输出表现变化情况。 使用绝对位置编码具有外推性上的限制,旋转位置编码则…

Ps:焦点堆栈

焦点堆栈 Focus Stacking是一种摄影和图像处理技术,通过合并多张在不同焦距拍摄的照片来创建一张具有更大景深的图像,特别适用于微距摄影、风景摄影和任何需要在整个场景中保持尖锐对焦的情况。 ◆ ◆ ◆ 拍摄注意事项 1、使用三脚架 为了确保图像之间…

《小强升职记:时间管理故事书》阅读笔记

目录 前言 一、你的时间都去哪儿了 1.1 你真的很忙吗 1.2 如何记录和分析时间日志 1.3 如何找到自己的价值观 二、无压工作法 2.1 传说中的“四象限法则 2.2 衣柜整理法 三、行动时遇到问题怎么办? 3.1 臣服与拖延 3.2 如何做到要事第一? 3.…

【一周年】我的创作纪念日

今天,是我成为创作者的第366天,不知不觉,来CSDN已经一年啦~ 在这个特殊的日子,也给大家讲讲我的创作故事。 一、机缘 起初,刚认识CSDN时,我的高中生涯刚结束,顺利从一名懵懂的高中生变身为一名懵…

Linux第53步_移植ST公司的linux内核第5步_系统镜像打包并烧录到EMMC

本节主要学习系统镜像打包,然后将打包文件烧录到EMMC测试。 1、创建bootfs文件夹 1)、打开第1个终端 输入“ls回车” 输入“cd linux/回车”,切换到“linux”目录 输入“ls回车”,列出“linux”目录下的文件和文件夹 输入“cd atk-mp1/…

Go语言中的加密艺术:深入解析crypto/subtle库

Go语言中的加密艺术:深入解析crypto/subtle库 引言crypto/subtle库概览ConstantTimeCompare函数深入解析ConstantTimeSelect函数应用详解ConstantTimeLessOrEq函数实践指南安全编程实践性能优化与最佳实践与其他加密库的比较总结 引言 在当今快速发展的互联网时代&…

localStorage、sessionStorage、cookie区别

localStorage: localStorage 的生命周期是永久的,关闭页面或浏览器之后 localStorage 中的数据也不会消失。localStorage 除非主动删除数据,否则数据永远不会消失 sessionStorage: sessionStorage 的生命周期是仅在当前会话下有效。sessionStorage 引入…

C语言第二十四弹---指针(八)

✨个人主页: 熬夜学编程的小林 💗系列专栏: 【C语言详解】 【数据结构详解】 指针 1、数组和指针笔试题解析 1.1、字符数组 1.1.1、代码1: 1.1.2、代码2: 1.1.3、代码3: 1.1.4、代码4: 1…