triton 模型部署

Triton Inference server是Nvidia开源的模型部署服务的框架。

源码地址：https://github.com/triton-inference-server

用来加速GPU推理。

triton支持tensorrt,tensorflow,pytorch等模型后端。

部署流程:

1.docker拉镜像

docker pull nvcr.io/nvidia/tritonserver:<xx.yy>-py3

xx.yy对应你想要的版本，我的是21.07

2.构造模型文件

模型文件目录结构：

  <model-repository-path>/
    <model-name>/
      [config.pbtxt]
      [<output-labels-file> ...]
      <version>/
        <model-definition-file>
      <version>/
        <model-definition-file>

tensorflow的模型文件目录结构：

  <model-repository-path>/
    <model-name>/
      config.pbtxt
      1/
        model.graphdef

或者是savedModel格式：

  <model-repository-path>/
    <model-name>/
      config.pbtxt
      1/
        model.savedmodel/
           <saved-model files>

pytorch的就是。model.pt

  <model-repository-path>/
    <model-name>/
      config.pbtxt
      1/
        model.pt

所以主要包括两个方面。一个config.pbtxt,一个model文件。

首先将tensorflow或pytorch模型的输入输出，写成config.pbtxt文件。

platform是要有的，可以是tensorrt,tensorflow,pytorch等。

max_batch_size是最大batch数，

dynamic_batch是动态batch,可以将某个时间段内的请求做成一个batch输入服务器，

max_deque_delay_microseconds:最大微秒队列等待，做成batch的最大等待的时间。

除此之外，config还可以reshape tensor，instance group对GPU分组请求，Sequence Batcher等功能。

  platform: "tensorflow"
  max_batch_size: 8
  input [
    {
      name: "input0"
      data_type: TYPE_FP32
      dims: [ 256 ]
    }
  ]
  output [
    {
      name: "output0"
      data_type: TYPE_FP32
      dims: [ 256,16 ]
    }
  ]
  dynamic_batching {
    preferred_batch_size: [ 4, 8 ]
    max_queue_delay_microseconds: 100
  }

2.做model文件。

要将model文件转换成需要的格式。

torch版本。

import os

os.environ['no_cuda'] = 'False'
os.environ['use_tensorrt'] = 'True'

import torch
from mybert import TokenClassification
import pickle


class WrappedModel(torch.nn.Module):
    def __init__(self):
        super(WrappedModel, self).__init__()
        self.model = TokenClassification.from_pretrained(
            '/path/to/model').cuda()
        # self.model = BertForSequenceClassification.from_pretrained('bert-base-uncased').cuda()

    def forward(self, inp0, inp1, inp2, inp3):
        return self.model(input_ids=inp0, attention_mask=inp1, token_type_ids=inp2, position_ids=inp3)




x = torch.zeros((4, 254), dtype=torch.long).cuda()  # bsz , seqlen
pt_model = WrappedModel().eval()
traced_script_module = torch.jit.trace(pt_model, (x[0], x[1], x[2], x[3]))
traced_script_module.save("model.pt")

tensorflow

import tensorflow as tf
from transformers import *
import numpy as np
tf.config.optimizer.set_jit(True)



class WrappedModel(tf.Module):
    def __init__(self):
        super(WrappedModel, self).__init__()
        self.model = tf.keras.models.load_model(r'/path/to/model')

    @tf.function
    def __call__(self, x):
        predict = self.model(x)
        return predict 


model = WrappedModel()
# example = tf.zeros((4,128), dtype=tf.int32) # bsz x seqlen
call = model.__call__.get_concrete_function(tf.TensorSpec([None, None, None], tf.int32, name='input_0'))  # bsz x seqlen
tf.saved_model.save(model, "model.savedmodel", signatures=call)

3.模型转换完成之后，直接运行即可。

GPU版本：

docker run --gpus all --rm -p11000:8000 -p11001:8001 -p11002:8002 -v /path/to/model:/models nvcr.io/nvidia/tritonserver:21.07-py3 tritonserver --model-repository=/models

CPU版本：去掉--gpus all 选项即可

docker run --rm -p11000:8000 -p11001:8001 -p11002:8002 -v /path/to/model:/models nvcr.io/nvidia/tritonserver:21.07-py3 tritonserver --model-repository=/models

最后运行成功，会显示，model ready.