先介绍TensorRT、Triton的关系和区别：
TensorRT：为inference（推理）为生，是NVIDIA研发的一款针对深度学习模型在GPU上的计算，显著提高GPU上的模型推理性能。即一种专门针对高性能推理的模型框架，也可以解析其他框架的模型如tensorflow、torch。
主要优化手段如下：
Triton：类似于TensorFlow Serving，但triton包括server和client。
triton serving能够实现不同模型的统一部署和服务，提供http和grpc协议，给triton client请求模型推理。
---------------------------------------分割线------------------------------------------------
如果是要将模型和推理嵌入在服务或软硬件中，那么TensorRT是很好的选择，使用它来加载模型进行推理，提升性能（tensorrt runtime）；
不然，常规的做法是模型推理和其他业务隔离，模型统一部署在triton server，然后其他业务通过triton client来进行模型推理的请求。
实验环境：Ubuntu18.04， GeForce RTX 2080Ti
Triton部署
安装
通过docker的形式，首先拉取镜像
# <xx.yy>为Triton的版本 docker pull nvcr.io/nvidia/tritonserver:<xx.yy>-py3# 例如，拉取 20.12 docker pull nvcr.io/nvidia/tritonserver:20.12-py3
<要注意不同版本的tritonserver对cuda驱动最低版本要求，以及对应的tensorrt版本>
例如，20.12的版本需要NVIDIA Driver需要455以上，支持TensorRT 7.2.2。TensorRT版本要对应，不然模型可能会无法部署。
其他版本信息可以前往官网查看：https://docs.nvidia.com/deeplea
启动
CPU版本的启动 NVIDIA Triton Inference Server
docker run --rm -p8000:8000 -p8001:8001 -p8002:8002 -v/full/path/to/docs/examples/model_repository:/models nvcr.io/nvidia/tritonserver:22.01-py3 tritonserver --model-repository=/models
GPU版本的启动，使用1个gpu
docker run --gpus=1 --rm -p8000:8000 -p8001:8001 -p8002:8002 -v/full/path/to/docs/examples/model_repository:/models nvcr.io/nvidia/tritonserver:22.01-py3 tritonserver --model-repository=/models
- /full/path/to/docs/examples/model_repository：模型仓库的路径。除了本地文件系统，还支持Google Cloud、S3、Azure这些云存储：https://github.com/triton-inference-server/server/blob/main/docs/model_repository.md
- --rm：表示容器停止运行时会删除容器
--gpus=1: 分配 1 个 GPU 资源给容器使用。- 8000为http端口，8001为grpc端口
-p8000:8000 -p8001:8001 -p8002:8002: 将容器内部的 8000、8001 和 8002 端口映射到宿主机的对应端口。这样可以从宿主机访问容器内部的服务。-v/full/path/to/docs/examples/model_repository:/models: 将宿主机上的/full/path/to/docs/examples/model_repository目录挂载到容器内的/models目录。这样容器可以访问宿主机上的模型文件。nvcr.io/nvidia/tritonserver:22.01-py3: 使用 NVIDIA 提供的 Triton Inference Server 22.01 版本的 Python 3 镜像作为容器的基础镜像。tritonserver --model-repository=/models: 启动 Triton Inference Server 服务,并指定模型仓库目录为/models，也就是我们挂载的宿主机目录。
正常启动的话，可以看到部署的模型运行状态，以及对外提供的服务端口
模型生成
Triton支持以下模型：TensorRT、ONNX、TensorFlow、Torch、OpenVINO、DALI，还有Python backend自定义生成的Python模型。
我们以一个简单的模型结构来演示：
我们以一个简单的模型结构来演示：
- INPUT0节点通过四则运算得到OUTPUT0节点；
- INPUT1节点通过embedding table映射为OUTPUT1
TensorFlow
tensorflow可以生成SavedModel或者GraphDef的模型格式
SavedModel模型需要按照以下的目录结构进行存储：
<model-repository-path>/<model-name>/config.pbtxt1/model.savedmodel/<saved-model files>
GraphDef：
<model-repository-path>/<model-name>/config.pbtxt1/model.graphdef
import os import tensorflow as tf from tensorflow.python.framework import graph_iodef create_modelfile(model_version_dir, max_batch,save_type="graphdef",version_policy=None):# your model netinput0_shape = [None, 2]input1_shape = [None, 2]x1 = tf.placeholder(tf.float32, input0_shape, name='INPUT0')inputs_id = tf.placeholder(tf.int32, input1_shape, name='INPUT1')out = tf.add(tf.multiply(x1, 0.5), 2)embedding = tf.get_variable("embedding_table", shape=[100, 10])pre = tf.nn.embedding_lookup(embedding, inputs_id)out0 = tf.identity(out, "OUTPUT0")out1 = tf.identity(pre, "OUTPUT1")try:os.makedirs(model_version_dir)except OSError as ex:pass # ignore existing dirwith tf.Session() as sess:sess.run(tf.global_variables_initializer())if save_type == 'graphdef':create_graphdef_modelfile(model_version_dir, sess,outputs=["OUTPUT0", "OUTPUT1"])elif save_type == 'savemodel':create_savedmodel_modelfile(model_version_dir,sess,inputs={"INPUT0": x1,"INPUT1": inputs_id},outputs={"OUTPUT0": out,"OUTPUT1": pre})else:raise ValueError("save_type must be one of ['tensorflow_graphdef', 'tensorflow_savedmodel']")create_modelconfig(models_dir=os.path.dirname(model_version_dir),max_batch=max_batch,save_type=save_type,version_policy=version_policy)def create_graphdef_modelfile(model_version_dir, sess, outputs):""" tensorflow graphdef只能保存constant，无法保存Variable可以借助tf.graph_util.convert_variables_to_constants将Variable转化为constant:param model_version_dir::param sess::return:""" graph = sess.graph.as_graph_def()new_graph = tf.graph_util.convert_variables_to_constants(sess=sess,input_graph_def=graph,output_node_names=outputs)graph_io.write_graph(new_graph,model_version_dir,"model.graphdef",as_text=False)def create_savedmodel_modelfile(model_version_dir, sess, inputs, outputs):""" :param model_version_dir::param sess::param inputs: dict, {input_name: input_tensor}:param outputs: dict, {output_name: output_tensor}:return:""" tf.saved_model.simple_save(sess,model_version_dir + "/model.savedmodel",inputs=inputs,outputs=outputs)
torch
pytorch模型的目录结构格式：<model-repository-path>/ <model-name>/ config.pbtxt 1/ model.pt
import os import torch from torch import nnclass MyNet(nn.Module):def __init__(self):super(MyNet, self).__init__()self.embedding = nn.Embedding(num_embeddings=100,embedding_dim=10)def forward(self, input0, input1):# tf.add(tf.multiply(x1, 0.5), 2)output0 = torch.add(torch.multiply(input0, 0.5), 2)output1 = self.embedding(input1)return output0, output1def create_modelfile(model_version_dir, max_batch,version_policy=None):# your model net# 定义输入的格式example_input0 = torch.zeros([2], dtype=torch.float32)example_input1 = torch.zeros([2], dtype=torch.int32)my_model = MyNet()traced = torch.jit.trace(my_model, (example_input0, example_input1))try:os.makedirs(model_version_dir)except OSError as ex:pass # ignore existing dirtraced.save(model_version_dir + "/model.pt")
ONNX
ONNX的目录结构：
<model-repository-path>/<model-name>/config.pbtxt1/model.onnx
ONNX提供一种开源的深度学习和传统的机器学习模型格式，目的在于模型在不同框架之间进行转移。
下面我们介绍最常用的tensorflow和torch模型转成ONNX的方法。
tensorflow模型 --> ONNX
pip install -U tf2onnx# savedmodel python -m tf2onnx.convert --saved-model tensorflow-model-path --output model.onnx# checkpoint python -m tf2onnx.convert --checkpoint tensorflow-model-meta-file-path --output model.onnx --inputs input0:0,input1:0 --outputs output0:0# graphdef python -m tf2onnx.convert --graphdef tensorflow-model-graphdef-file --output model.onnx --inputs input0:0,input1:0 --outputs output0:0
torch --> ONNX
import os import torch import torch.onnxdef torch2onnx(model_version_dir, max_batch):# 定义输入的格式example_input0 = torch.zeros([max_batch, 2], dtype=torch.float32)example_input1 = torch.zeros([max_batch, 2], dtype=torch.int32)my_model = MyNet()try:os.makedirs(model_version_dir)except OSError as ex:pass # ignore existing dirtorch.onnx.export(my_model,(example_input0, example_input1),os.path.join(model_version_dir, 'model.onnx'),# 输入节点的名称input_names=("INPUT0", "INPUT1"),# 输出节点的名称output_names=("OUTPUT0", "OUTPUT1"),# 设置batch_size的维度dynamic_axes={"INPUT0": [0], "INPUT1": [0], "OUTPUT0": [0], "OUTPUT1": [0]},verbose=True)
TensorRT
需要注意：TensorRT仅支持GPU。
<model-repository-path>/<model-name>/config.pbtxt1/model.plan
比较推荐的方式是从ONNX解析得到TensorRT模型（TensorRT）
import tensorrt as trt import osdef onnx2trt(model_version_dir, onnx_model_file, max_batch):logger = trt.Logger(trt.Logger.WARNING)builder = trt.Builder(logger)# The EXPLICIT_BATCH flag is required in order to import models using the ONNX parsernetwork = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser = trt.OnnxParser(network, logger)success = parser.parse_from_file(onnx_model_file)for idx in range(parser.num_errors):print(parser.get_error(idx))if not success:pass # Error handling code hereprofile = builder.create_optimization_profile()# INPUT0可以接收[1, 2] -> [max_batch, 2]的维度profile.set_shape("INPUT0", [1, 2], [1, 2], [max_batch, 2])profile.set_shape("INPUT1", [1, 2], [1, 2], [max_batch, 2])config = builder.create_builder_config()config.add_optimization_profile(profile)# tensorrt 8.x# config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20) # 1 MiB# tensorrt 7.xconfig.max_workspace_size = 1 << 20try:engine_bytes = builder.build_serialized_network(network, config)except AttributeError:engine = builder.build_engine(network, config)engine_bytes = engine.serialize()del enginewith open(os.path.join(model_version_dir, 'model.plan'), "wb") as f:f.write(engine_bytes)
模型配置文件
name: "tf_savemodel" platform: "tensorflow_savedmodel" max_batch_size: 8 version_policy: { latest { num_versions: 1 }} input [{name: "INPUT0"data_type: TYPE_FP32dims: [ 2 ]},{name: "INPUT1"data_type: TYPE_INT32dims: [ 2 ]} ] output [{name: "OUTPUT0"data_type: TYPE_FP32dims: [ 2 ]},{name: "OUTPUT1"data_type: TYPE_FP32dims: [ 2,10 ]} ]
name：模型名称，要跟模型路径对应。
platform：不同的模型存储格式都有自己对应的值。
max_batch_size：最大的batch_size，客户端超过这个batch_size的请求会报错。
version_policy：版本控制，这里是使用最新的一个版本。
input、output：输入和输出节点的名称，数据类型，维度。
维度一般不包括batch_size这个维度；
下表为不同框架对应的platform：
下表是不同框架的数据类型对应关系：Model Config是配置文件的，API是triton client。其他框架是c++源码的命名空间，不过很好理解，主要包括16位和32位的int和float等等。
Triton Client
上述提到了，我们可以通过triton client来进行模型推理的请求，并且提供了http和grpc两种协议。
接下来，将以python来演示，仍然是上面那个简单的模型请求例子。
# 安装依赖包
pip install tritonclient[all]
import gevent.ssl import numpy as np import tritonclient.http as httpclientdef client_init(url="localhost:8000",ssl=False, key_file=None, cert_file=None, ca_certs=None, insecure=False,verbose=False):""" :param url::param ssl: Enable encrypted link to the server using HTTPS:param key_file: File holding client private key:param cert_file: File holding client certificate:param ca_certs: File holding ca certificate:param insecure: Use no peer verification in SSL communications. Use with caution:param verbose: Enable verbose output:return:""" if ssl:ssl_options = {}if key_file is not None:ssl_options['keyfile'] = key_fileif cert_file is not None:ssl_options['certfile'] = cert_fileif ca_certs is not None:ssl_options['ca_certs'] = ca_certsssl_context_factory = Noneif insecure:ssl_context_factory = gevent.ssl._create_unverified_contexttriton_client = httpclient.InferenceServerClient(url=url,verbose=verbose,ssl=True,ssl_options=ssl_options,insecure=insecure,ssl_context_factory=ssl_context_factory)else:triton_client = httpclient.InferenceServerClient(url=url, verbose=verbose)return triton_clientdef infer(triton_client, model_name,input0='INPUT0', input1='INPUT1',output0='OUTPUT0', output1='OUTPUT1',request_compression_algorithm=None,response_compression_algorithm=None):""" :param triton_client::param model_name::param input0::param input1::param output0::param output1::param request_compression_algorithm: Optional HTTP compression algorithm to use for the request body on client side.Currently supports "deflate", "gzip" and None. By default, no compression is used.:param response_compression_algorithm::return:""" inputs = []outputs = []# batch_size=8# 如果batch_size超过配置文件的max_batch_size，infer则会报错# INPUT0、INPUT1为配置文件中的输入节点名称inputs.append(httpclient.InferInput(input0, [8, 2], "FP32"))inputs.append(httpclient.InferInput(input1, [8, 2], "INT32"))# Initialize the data# np.random.seed(2022)inputs[0].set_data_from_numpy(np.random.random([8, 2]).astype(np.float32), binary_data=False)# np.random.seed(2022)inputs[1].set_data_from_numpy(np.random.randint(0, 20, [8, 2]).astype(np.int32), binary_data=False)# OUTPUT0、OUTPUT1为配置文件中的输出节点名称outputs.append(httpclient.InferRequestedOutput(output0, binary_data=False))outputs.append(httpclient.InferRequestedOutput(output1,binary_data=False))query_params = {'test_1': 1, 'test_2': 2}results = triton_client.infer(model_name=model_name,inputs=inputs,outputs=outputs,request_compression_algorithm=request_compression_algorithm,response_compression_algorithm=response_compression_algorithm)print(results)# 转化为numpy格式print(results.as_numpy(output0))print(results.as_numpy(output1))
grpc的代码基本相同，就不展示了，
