前言
本篇介绍以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 { max_queue_delay_microseconds: 2000 }该文件决定了模型的输入输出的维度，服务策略等内容，重点内容如下
- backend：推理后端，本例采用Python实现的自定义客户端，在Python中使用了tensorrt的API，因此本质上是tensorrt的后端
- max_batch_size：一次推理的最大批次，超过该值会报错，max_batch_size通常和动态批处理dynamic_batching一起使用，max_batch_size会作为停止合并的一个条件
- input/output：输入和输入的定义，变量名自定义，但是必须和Python后端脚本一致，dims代表维度，-1代表不定长，data_type代表类型，具体使用请参考Triton的教程
- instance_group：多实例设置，kind代表设备，KIND_GPU为GPU设备，也可以指定CPU，gpus指定GPU设备号，多个id就是指定多gpu，count代表实例数，具体是每个GPU/CPU下的实例数，本例中代表0号GPU启2个实例
- dynamic_batching：动态批处理，服务端会自动合并请求，从而尽量以批量推理的方式来代替单条请求推理，提高吞吐量，因此服务端会主动等待max_queue_delay_microseconds时间，在这段时间内服务端会将所有请求合并，合并完成后再输送给推理后端，推理完成后合并的结果会再拆成单条请求的形式，因此对客户端无感。默认情况下如果不设置dynamic_batching，Triton不会进行动态批处理
自定义Python客户端需要在版本号文件夹下设置一个model.py文件，该文件内部实现了后端推理逻辑，work目录为服务运行过程中自动生成，不需要理会。
Triton服务端代码实现
服务端代码实现在model.py中，具体的在其中实现trt文件的读取，客户端数据的获取，模型推理，响应返回，本例如下
import os # 设置显存空闲block最大分割阈值 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32' # 设置work目录 os.environ['TRANSFORMERS_CACHE'] = os.path.dirname(os.path.abspath(__file__)) + "/work/" os.environ['HF_MODULES_CACHE'] = os.path.dirname(os.path.abspath(__file__)) + "/work/" import json import triton_python_backend_utils as pb_utils import sys import gc import time import logging from transformers import BertTokenizer import tensorrt as trt import numpy as np import torch import trtutils as trtu gc.collect() logging.basicConfig(format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s', level=logging.INFO) TRT_LOGGER = trt.Logger(trt.Logger.WARNING) class TritonPythonModel: def initialize(self, args): # You must parse model_config. JSON string is not parsed here self.model_config = json.loads(args['model_config']) output_config = pb_utils.get_output_config_by_name(self.model_config, "embeddings") # Convert Triton types to numpy types self.output_response_dtype = pb_utils.triton_string_to_numpy(output_config['data_type']) # trt engine 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 = os.path.dirname(os.path.abspath(__file__)) + "/m3e-base-trt" self.engine = get_engine(engine_model_path + "/m3e-base.trt") self.context = self.engine.create_execution_context() self.context.active_optimization_profile = 0 self.tokenizer = BertTokenizer.from_pretrained(engine_model_path + "/vocab.txt") logging.info("model init success") def execute(self, requests): responses = [] # TODO 记录下每个请求的数据和数据batch大小 batch_text, batch_len = [], [] for request in requests: text = pb_utils.get_input_tensor_by_name(request, "text").as_numpy().astype("S") text = np.char.decode(text, "utf-8").squeeze(1).tolist() batch_text.extend(text) batch_len.append(len(text)) # 日志输出传入信息 in_log_info = { "text": batch_text, } logging.info(in_log_info) # tokenizer encode = self.tokenizer.batch_encode_plus(batch_text, max_length=512, truncation=True, padding=True) input_ids, attention_mask = np.array(encode["input_ids"]).astype(np.int32), np.array( encode["attention_mask"]).astype( np.int32) origin_input_shape = self.context.get_binding_shape(0) origin_input_shape[0], origin_input_shape[1] = input_ids.shape self.context.set_binding_shape(0, origin_input_shape) self.context.set_binding_shape(1, origin_input_shape) inputs, outputs, bindings, stream = trtu.allocate_buffers_v2(self.engine, self.context) inputs[0].host = input_ids inputs[1].host = attention_mask trt_outputs = trtu.do_inference_v2(self.context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream) token_embeddings = trt_outputs[0].reshape(input_ids.shape[0], input_ids.shape[1], 768) # mean pool attention_mask = torch.LongTensor(attention_mask) token_embeddings = torch.tensor(token_embeddings) input_mask_expanded = ( attention_mask.unsqueeze(-1).expand(token_embeddings.size()).to(token_embeddings.dtype) ) sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask = input_mask_expanded.sum(1) sum_mask = torch.clamp(sum_mask, min=1e-9) token_embeddings = sum_embeddings / sum_mask # 归一化 token_embeddings = torch.nn.functional.normalize(token_embeddings, p=2, dim=1) token_embeddings = token_embeddings.detach().cpu().numpy().tolist() # 日志输出处理后的信息 out_log_info = { "embeddings": token_embeddings } #logging.info(out_log_info) # TODO 响应数要和请求数一致 start = 0 for i in range(len(requests)): end = start + batch_len[i] out_tensor = pb_utils.Tensor("embeddings", np.array(token_embeddings[start:end]).astype(self.output_response_dtype)) start += batch_len[i] final_inference_response = pb_utils.InferenceResponse(output_tensors=[out_tensor]) responses.append(final_inference_response) return responses def finalize(self): print('Cleaning up...')其中tensorrt只支持int32，因此需要手动将numpy数据类型转化为int32否则推理报错，trtu.do_inference_v2完成了模型推理。注意从客户端拿到的text和返回的embeddings命名都是要和config.pbtxt保持一致的。
trtutils为现成的推理工具方法，直接使用即可，代码如下
import argparse import os import numpy as np import pycuda.autoinit import pycuda.driver as cuda import tensorrt as trt try: # Sometimes python does not understand FileNotFoundError FileNotFoundError except NameError: FileNotFoundError = IOError EXPLICIT_BATCH = 1
