前言

本篇介绍以Triton作为推理服务器，TensorRT作为推理后端，部署句嵌入向量模型m3e-base的工程方案和实现，句嵌入模型本质上是Bert结构，本案例可以推广到更一般的深度学习模型部署场景。

内容摘要

推理服务器和推理后端介绍TensorRT+Triton环境搭建Bert模型转化为ONNX中间表示ONNX中间表示编译为TensorRT模型文件Triton服务端参数配置Triton服务端代码实现Triton服务端启动HTTP客户端请求TensorRT前后压测结果对比

推理服务器和推理后端介绍

在Triton+TensorRT的组合中，Triton是<code>推理服务器，TensorRT是推理后端，两者都是NVIDIA推出的推理部署服务组件，Triton原名TensorRT Inference Server，是专供于TensorRT后端的推理服务器，由于TensorRT Inference Server支持的后端越来越多，因此其改名为Triton，Triton+TensorRT一直是NVIDIA主推的部署方式。

模型的推理服务由两块组件构成，分别是推理服务器和推理后端，推理服务器负责处理客户端的请求输送给推理后端，推理后端负责输出推理结果，经过推理服务器返回给客户端，两者的工作示意图如下

推理服务器和推理后端示意图

推理服务的优化需要推理服务器和推理后端的共同优化，其中推理后端的优化主要体现在推理引擎<code>内核的优化使得推理的性能更高，延迟更低，推理服务器的优化主要体现在推理策略、调度策略的优化，通过辅助策略来协助后端更好地推理，提高吞吐量。

推理服务器接受客户端请求，并为后端推理提供必要的前提准备，包括

前处理：对请求的数据进行预先处理，使得服务后端模型的要求，例如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。

<code>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官网下载

解压

<code>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',code>

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 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def GiB(val):

return val * 1 << 30

def add_help(description):

parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)

args, _ = parser.parse_known_args()

def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""):code>

'''

Parses sample arguments.

Args:

description (str): Description of the sample.

subfolder (str): The subfolder containing data relevant to this sample

find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.

Returns:

str: Path of data directory.

'''

# Standard command-line arguments for all samples.

kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")

parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument("-d", "--datadir",

help="Location of the TensorRT sample data directory, and any additional data directories.",code>

action="append", default=[kDEFAULT_DATA_ROOT])code>

args, _ = parser.parse_known_args()

def get_data_path(data_dir):

# If the subfolder exists, append it to the path, otherwise use the provided path as-is.

data_path = os.path.join(data_dir, subfolder)

if not os.path.exists(data_path):

if data_dir != kDEFAULT_DATA_ROOT:

print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")

data_path = data_dir

# Make sure data directory exists.

if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:

print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(

data_path))

return data_path

data_paths = [get_data_path(data_dir) for data_dir in args.datadir]

return data_paths, locate_files(data_paths, find_files, err_msg)

def locate_files(data_paths, filenames, err_msg=""):code>

"""

Locates the specified files in the specified data directories.

If a file exists in multiple data directories, the first directory is used.

Args:

data_paths (List[str]): The data directories.

filename (List[str]): The names of the files to find.

Returns:

List[str]: The absolute paths of the files.

Raises:

FileNotFoundError if a file could not be located.

"""

found_files = [None] * len(filenames)

for data_path in data_paths:

# Find all requested files.

for index, (found, filename) in enumerate(zip(found_files, filenames)):

if not found:

file_path = os.path.abspath(os.path.join(data_path, filename))

if os.path.exists(file_path):

found_files[index] = file_path

# Check that all files were found

for f, filename in zip(found_files, filenames):

if not f or not os.path.exists(f):

raise FileNotFoundError(

"Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg))

return found_files

# Simple helper data class that's a little nicer to use than a 2-tuple.

class HostDeviceMem(object):

def __init__(self, host_mem, device_mem):

self.host = host_mem

self.device = device_mem

def __str__(self):

return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

def __repr__(self):

return self.__str__()

# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.

def allocate_buffers(engine):

inputs = []

outputs = []

bindings = []

stream = cuda.Stream()

for binding in engine:

size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size # max_batch_size=1

dtype = trt.nptype(engine.get_binding_dtype(binding))

# Allocate host and device buffers

host_mem = cuda.pagelocked_empty(size, dtype)

device_mem = cuda.mem_alloc(host_mem.nbytes) # # nbytes表示数组中的所有数据消耗掉的字节数

# Append the device buffer to device bindings.

bindings.append(int(device_mem))

# Append to the appropriate list.

if engine.binding_is_input(binding):

inputs.append(HostDeviceMem(host_mem, device_mem))

else:

outputs.append(HostDeviceMem(host_mem, device_mem))

return inputs, outputs, bindings, stream

# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.

def allocate_buffers_v2(engine, context):

"""

Allocates host and device buffer for TRT engine inference.

This function is similiar to the one in ../../common.py, but

converts network outputs (which are np.float32) appropriately

before writing them to Python buffer. This is needed, since

TensorRT plugins doesn't support output type description, and

in our particular case, we use NMS plugin as network output.

Args:

engine (trt.ICudaEngine): TensorRT engine

Returns:

inputs [HostDeviceMem]: engine input memory

outputs [HostDeviceMem]: engine output memory

bindings [int]: buffer to device bindings

stream (cuda.Stream): cuda stream for engine inference synchronization

"""

inputs = []

outputs = []

bindings = []

stream = cuda.Stream()

for i, binding in enumerate(engine):

# binding:input_ids,input_mask,output

# print(context.get_binding_shape(i)) # (input_ids,input_mask,output).shape (1,105)

size = trt.volume(context.get_binding_shape(i)) # 1*105

# dims = engine.get_binding_shape(binding)

# if dims[1] < 0:

# size *= -1

dtype = trt.nptype(engine.get_binding_dtype(binding)) # DataType.FLOAT

# print(dtype) # <class 'numpy.float32'>

# Allocate host and device buffers

host_mem = cuda.pagelocked_empty(size, dtype)

device_mem = cuda.mem_alloc(host_mem.nbytes)

# Append the device buffer to device bindings.

bindings.append(int(device_mem))

# Append to the appropriate list.

if engine.binding_is_input(binding):

inputs.append(HostDeviceMem(host_mem, device_mem))

else:

outputs.append(HostDeviceMem(host_mem, device_mem))

return inputs, outputs, bindings, stream

# This function is generalized for multiple inputs/outputs.

# inputs and outputs are expected to be lists of HostDeviceMem objects.

def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):

# Transfer input data to the GPU.

[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]

# Run inference. batch_size = 1

context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)

# Transfer predictions back from the GPU.

[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]

# Synchronize the stream

stream.synchronize()

# Return only the host outputs.

return [out.host for out in outputs]

# This function is generalized for multiple inputs/outputs for full dimension networks.

# inputs and outputs are expected to be lists of HostDeviceMem objects.

def do_inference_v2(context, bindings, inputs, outputs, stream):

# Transfer input data to the GPU.

[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]

# Run inference.

context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)

# Transfer predictions back from the GPU.

[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]

# Synchronize the stream

stream.synchronize()

# Return only the host outputs.

return [out.host for out in outputs]

由于m3e模型的特殊性，需要对模型推理结果做后处理，本例增加了mean pool和归一化操作，代码摘自sentence_transformers的源码，由于在模型转换阶段只对原生的Bert模型进行了转化，因此后处理需要额外补充进来。

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【<code>保证100%免费】