建立 Triton Server

Triton Server 是 NVIDIA 推出的一个用于机器学习模型布置的开源渠道，支持 TensorFlow、PyTorch、ONNX 等多种模型格式。下面是建立 Triton Server 的步骤：

  1. 装置 Docker

Triton Server 是基于 Docker 运转的，因此需要先装置 Docker。

  1. 下载 Triton Server 镜像

在终端中履行以下指令下载 Triton Server 镜像：

docker pull nvcr.io/nvidia/tritonserver:21.09-py3

其间 21.09 是 Triton Server 的版本号，py3 表明运用 Python 3.x。

  1. 发动 Triton Server

在终端中履行以下指令发动 Triton Server：

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

其间 /path/to/model/repository 是寄存模型的目录，能够自己指定。发动 Triton Server 后，能够经过 http://localhost:8000/v2/models 访问 Triton Server 的 API。

导出 PyTorch BERT 模型为 ONNX

  1. 装置 PyTorch 和 ONNX

在终端中履行以下指令装置 PyTorch 和 ONNX：

pip install torch==1.9.0
pip install onnx==1.9.0
  1. 导出 PyTorch BERT 模型为 ONNX

在 Python 脚本中履行以下代码将 PyTorch BERT 模型导出为 ONNX：

import torch
import onnx
# 加载 PyTorch BERT 模型
model = torch.load('/path/to/pytorch/bert/model')
# 转换为 ONNX 格式
input_names = ['input_ids', 'attention_mask', 'token_type_ids']
output_names = ['output']
dummy_input = (torch.zeros(1, 128), torch.zeros(1, 128), torch.zeros(1, 128))
onnx_model = onnx.export(model, dummy_input, '/path/to/onnx/bert/model', input_names=input_names, output_names=output_names)

其间 /path/to/pytorch/bert/model 是 PyTorch BERT 模型的途径，/path/to/onnx/bert/model 是导出的 ONNX 模型的途径。

布置 PyTorch BERT 模型到 Triton Server

  1. 创立模型配置文件

在 /path/to/model/repository 目录下创立一个名为 bert 的目录，并在该目录下创立一个名为 config.pbtxt 的文件，文件内容如下：

name: "bert"
platform: "pytorch_libtorch"
max_batch_size: 1
input [
  {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [ -1, 128 ]
  },
  {
    name: "attention_mask"
    data_type: TYPE_INT64
    dims: [ -1, 128 ]
  },
  {
    name: "token_type_ids"
    data_type: TYPE_INT64
    dims: [ -1, 128 ]
  }
]
output [
  {
    name: "output"
    data_type: TYPE_FP32
    dims: [ -1, 768 ]
  }
]

其间 name 是模型的名称，platform 是模型的布置渠道，max_batch_size 是最大批次大小，input 是模型的输入，output 是模型的输出。

  1. 仿制模型文件

将导出的 ONNX 模型文件仿制到 /path/to/model/repository/bert/1/ 目录下，并将文件名改为 model.pt。

  1. 发动 Triton Server

在终端中履行以下指令发动 Triton Server：

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

其间 /path/to/model/repository 是寄存模型的目录，能够自己指定。发动 Triton Server 后，能够经过 http://localhost:8000/v2/models 访问 Triton Server 的 API。

恳求模型的 HTTP 和 gRPC 代码示例

下面是恳求 PyTorch BERT 模型的 HTTP 和 gRPC 代码示例：

HTTP：

import requests
import json
url = 'http://localhost:8000/v2/models/bert/infer'
headers = {'Content-Type': 'application/json'}
data = {
    'inputs': [
        {
            'name': 'input_ids',
            'shape': [1, 128],
            'datatype': 'INT64',
            'data': [0] * 128
        },
        {
            'name': 'attention_mask',
            'shape': [1, 128],
            'datatype': 'INT64',
            'data': [0] * 128
        },
        {
            'name': 'token_type_ids',
            'shape': [1, 128],
            'datatype': 'INT64',
            'data': [0] * 128
        }
    ],
    'outputs': [
        {
            'name': 'output',
            'shape': [1, 768],
            'datatype': 'FP32'
        }
    ]
}
response = requests.post(url, headers=headers, data=json.dumps(data))
result = json.loads(response.content.decode('utf-8'))
print(result)

gRPC：

import grpc
import numpy as np
import onnxruntime_protobuf as ort_pb2
import onnxruntime_protobuf as ort_pb2_grpc
channel = grpc.insecure_channel('localhost:8001')
stub = ort_pb2_grpc.InferenceServerStub(channel)
inputs = [
    {
        'name': 'input_ids',
        'shape': [1, 128],
        'datatype': 'INT64',
        'data': np.zeros((1, 128), dtype=np.int64)
    },
    {
        'name': 'attention_mask',
        'shape': [1, 128],
        'datatype': 'INT64',
        'data': np.zeros((1, 128), dtype=np.int64)
    },
    {
        'name': 'token_type_ids',
        'shape': [1, 128],
        'datatype': 'INT64',
        'data': np.zeros((1, 128), dtype=np.int64)
    }
]
outputs = [
    {
        'name': 'output',
        'shape': [1, 768],
        'datatype': 'FP32'
    }
]
request = ort_pb2.InferRequest()
request.model_name = 'bert'
request.inputs.extend([ort_pb2.NamedTensor(name=input['name'], dims=input['shape'], data_type=input['datatype'], raw_data=input['data'].tobytes()) for input in inputs])
request.outputs.extend([ort_pb2.NamedTensor(name=output['name'], dims=output['shape'], data_type=output['datatype']) for output in outputs])
response = stub.Infer(request)
result = {}
for output in response.outputs:
    result[output.name] = np.frombuffer(output.raw_data, dtype=np.float32).reshape(output.dims)
print(result)

以上代码示例中，inputs 是模型的输入，outputs 是模型的输出。在恳求模型时，需要将输入数据的类型、形状和数据传递给 Triton Server，Triton Server 会返回输出数据的类型、形状和数据。

总结

本文介绍了如何运用 Triton Server 建立一个 PyTorch BERT 模型的推理服务，并供给了 HTTP 和 gRPC 恳求代码示例。经过运用 Triton Server，能够方便地进行模型推理服务的布置和办理，同时供给高效的推理服务。如果您想了解更多关于 Triton Server 的信息，能够参阅官方文档：docs.nvidia.com/deeplearnin…