前言
本篇介绍以Triton作为推理服务器,TensorRT作为推理后端,部署句嵌入向量模型m3e-base的工程方案和实现,句嵌入模型本质上是Bert结构,本案例可以推广到更一般的深度学习模型部署场景。
内容摘要
- 推理服务器和推理后端介绍
- TensorRT+Triton环境搭建
- Bert模型转化为ONNX中间表示
- ONNX中间表示编译为TensorRT模型文件
- Triton服务端参数配置
- Triton服务端代码实现
- Triton服务端启动
- HTTP客户端请求
- TensorRT前后压测结果对比
推理服务器和推理后端介绍
在Triton+TensorRT的组合中,Triton是推理服务器,TensorRT是推理后端,两者都是NVIDIA推出的推理部署服务组件,Triton原名TensorRT Inference Server,是专供于TensorRT后端的推理服务器,由于TensorRT Inference Server支持的后端越来越多,因此其改名为Triton,Triton+TensorRT一直是NVIDIA主推的部署方式。
模型的推理服务由两块组件构成,分别是推理服务器和推理后端,推理服务器负责处理客户端的请求输送给推理后端,推理后端负责输出推理结果,经过推理服务器返回给客户端,两者的工作示意图如下

推理服务器和推理后端示意图
推理服务的优化需要推理服务器和推理后端的共同优化,其中推理后端的优化主要体现在推理引擎内核的优化使得推理的性能更高,延迟更低,推理服务器的优化主要体现在推理策略、调度策略的优化,通过辅助策略来协助后端更好地推理,提高吞吐量。
推理服务器接受客户端请求,并为后端推理提供必要的前提准备,包括
- 前处理:对请求的数据进行预先处理,使得服务后端模型的要求,例如NLP领域,对输入的自然语言使用tokenizer进行分词编码
- 后处理:对后端模型的推理结果做再加工,符合客户端的需求,比如对输出的矩阵进行reshape转换,label映射等
- 模型预热:在服务完全启动之前,提供一些样例数据给到模型推理,使得模型完全初始化
- 动态批处理:服务端自动将请求时间接近的所携带的数据进行合并,从而批量推理提高GPU的利用率,等待请求的合并会带来一定的延迟
- 并发多实例:推理服务器设置多个实例来共同处理请求,分配资源来处理模型的负载,提高并发请求下的服务吞吐量
推理后端又分别推理表示和推理引擎,其内容如下
- 推理表示:即模型格式,是模型训练后输出的模型文件,例如Tensorflow的冻结图,PyTorch的bin格式
- 推理引擎:即支持该类模型格式的推理计算组件,Tensorflow和PyTorch这样的训练框架本身自带推理引擎,也有其他的更加优化的推理引擎,比如ONNXRuntime等。
业界常用的推理服务器和推理后端罗列如下

推理服务器和推理后端选型
本文要介绍的是以Triton作为推理服务器,以TensorRT作为推理后端的部署方案,其中Triton中的后端程序由Python实现,模型格式为TensorRT,使用Python后端下的TensorRT包实现对模型推理。
TensorRT+Triton环境搭建
笔者的环境为NVIDIA显卡驱动driver版本为535.154.05,cuda版本为12.2。下载Triton的Docker镜像,到NVIDIA查看符合cuda版本的镜像。 下载23.08版本的Triton镜像,对应的TensorRT版本为8.6.1.6,该镜像提供了推理服务器环境,是模型服务的基础镜像,该镜像的Python3版本为3.10。
docker pull nvcr.io/nvidia/tritonserver:23.08-py3
下载23.08版本的TensorRT镜像,该镜像的作用是使用trtexec将onnx模型格式转化为trt格式
docker pull nvcr.io/nvidia/tensorrt:23.08-py3
手动下载8.6.1.6版本的TensorRT,下载的目的是手动安装TensorRT的Python包,在推理的时候需要TensorRT的Python API实现推理

TensorRT官网下载
解压
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-12.0.tar.gz
下一步搭建基础镜像环境,需要在tritonserver镜像中安装如下Python包依赖,参考版本如下
torch 2.1.2+cu121
transformers 4.39.3
tensorrt 8.6.1
sentence-transformers 2.7.0
pycuda 2022.2.2
手动下载torch gpu版本和cuda,和TensorRT解压后的文件夹到pip_package下,一起映射到容器
# docker启动Triton镜像
docker run --rm -it -v /home/pip_package/:/home nvcr.io/nvidia/tritonserver:23.08-py3 /bin/bash
# 安装tensorrt
cd /home
pip install TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp310-none-linux_x86_64.whl
# 安装torch
pip install torch-2.1.2+cu121-cp310-cp310-linux_x86_64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple
# 安装transformers
pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
# 安装sentence-transformers
pip install sentence-transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
# 安装pycuda
pip install pycuda -i https://pypi.tuna.tsinghua.edu.cn/simple
将容器保存为一个新的镜像,至此环境搭建完毕
docker commit xxxxxxx triton_tensorrt_py_23.08:v1
Bert模型转化为ONNX中间表示
使用TensorRT作为后端推理模型必须将模型转化为trt格式,目前众多模型都支持转化为trt,但是支持程度层次不起,其中TensorRT对模型中间表示ONNX支持的最好,因此一般的做法是将tensorflow,pytorch的模型文件先转化为ONNX,再从ONNX转化为trt格式。
首先将m3e-base模型转化为ONNX格式,PyTorch API支持直接转化
from transformers import BertModel
model = BertModel.from_pretrained("./m3e-base").eval()
import torch
onnx_path = "./m3e-base.onnx"
input_ids = torch.LongTensor([[1, 2, 3], [2, 3, 4]])
attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 1]])
torch.onnx.export(model,
(input_ids, attention_mask),
onnx_path,
verbose=False,
opset_version=11,
input_names=['input_ids', 'attention_mask'],
output_names=['output'],
dynamic_axes={
"input_ids": {
0: "batch_size", 1: "max_seq_len"},
"attention_mask": {
0: "batch_size", 1: "max_seq_len"},
"output": {
0: "batch_size"}
})
其中input_names和output_names取名可以自定义,输入的顺序必须和模型forward顺序一致,dynamic_axes代表不定长的动态维度,指定维度索引和一个自定义命名,本例中input_ids,attention_mask的0,1维度都是不定长,output的0维度是不定长
ONNX中间表示编译为TensorRT模型文件
下一步将ONNX文件转化为trt格式,将m3e-base.onnx(/home/model_repository/目录下)映射到tensorrt容器内,使用trtexec进行转换,需要将宿主机的gpu挂入容器内
docker run --gpus=all --rm -it -v /home/model_repository/:/home nvcr.io/nvidia/tensorrt:23.08-py3 /bin/bash
trtexec --onnx=m3e-base.onnx \
--workspace=10000 \
--saveEngine=m3e-base.trt \
--minShapes=input_ids:1x1,attention_mask:1x1 \
--optShapes=input_ids:16x512,attention_mask:16x512 \
--maxShapes=input_ids:64x512,attention_mask:64x512
若日志显示PASSED代表转化成功,若显示空间不足报错请适当增大workspace,其中saveEngine代表模型输出的文件命中,minShapes,optShapes,maxShapes代表支持动态输入,指定最小尺寸和最大尺寸。转化完成后输出trt文件m3e-base.trt,将其映射到triton_tensorrt_py_23.08:v1容器中,测试是否能够正常被tensorrt的Python API读取
docker run --rm -it --gpus=all -v /home/model_repository:/home triton_tensorrt_py_23.08:v1 /bin/bash
用容器内的Python3来读取trt文件
root@a10830d0aeec:/home# python3
Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorrt as trt
>>> TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
>>> def get_engine(engine_file_path):
... print("Reading engine from file {}".format(engine_file_path))
... with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
... engine = runtime.deserialize_cuda_engine(f.read())
... return engine
...
>>> engine_model_path = "m3e-base.trt"
>>> engine = get_engine(engine_model_path)
Reading engine from file m3e-base.trt
>>>
若没有报错则表示trt模型转化成功
Triton服务端参数配置
Triton部署模型的服务端有严格的文件目录和文件形式要求,首先创建一个模型目录命名为m3e-base-trt,本例的文件格式如下
(base) [root@localhost m3e-base-trt]# tree
.
├── 1
│ ├── m3e-base-trt
│ │ ├── m3e.trt
│ │ └── vocab.txt
│ ├── model.py
│ ├── __pycache__
│ │ ├── model.cpython-310.pyc
│ │ ├── model.cpython-38.pyc
│ │ ├── trtutils.cpython-310.pyc
│ │ └── trtutils.cpython-38.pyc
│ ├── trtutils.py
│ └── work
│ └── version.txt
└── config.pbtxt
该目录下的1代表模型版本,可以取任意数字id作为文件名代表模型版本,默认情况下Triton以最大的那个模型版本作为服务。config.pbtxt为模型的服务端配置文件,配置如下
(base) [root@localhost m3e-base-trt]# cat config.pbtxt
name: "m3e-base-trt"
backend: "python"
max_batch_size: 32
input [
{
name: "text"
dims: [ -1 ]
data_type: TYPE_STRING
}
]
output [
{
name: "embeddings"
dims: [ 768 ]
data_type: TYPE_FP32
}
]
instance_group [
{
count: 2
kind: KIND_GPU
gpus: [ 0 ]
}
]
dynamic_batching {

最低0.47元/天 解锁文章
2327

被折叠的 条评论
为什么被折叠?



