NVIDIA Triton Inference Server

发表于 2024-06-09 更新于 2024-11-19 分类于工具/tool 阅读次数：

本文字数： 9.4k 阅读时长 ≈ 17 分钟

NVIDIA Triton 是一款由 NVIDIA 开发的高性能推理服务软件，不仅能够支持多种主流的推理框架，如 PyTorch、ONNX Runtime 和 TensorRT 等，还具备强大的多GPU负载均衡能力，确保高效利用计算资源。NVIDIA Triton 通过提供统一的接口来管理不同框架下的模型，能够大大提高模型部署的灵活性与效率。

概述

Quick Start

启动和使用NVIDIA Triton主要有3个步骤：

创建模型库
启动Triton Server
发送推理请求

创建模型库

创建文件夹model_repository，里面添加了yolov5s onnx模型，具体文件布局如下：

├── model_repository
│   └── det_yolov5s
│       ├── 1
│       │   └── model.onnx
│       └── config.pbtxt

配置yolov5s算法需要指定一个配置文件config.pbtxt和对应模型model.onnx（必须命名为model.<suffix>，比如onnx模型是model.onnx，tensorrt模型是model.plan）。配置文件config.pbtxt最小内容如下：

name: "det_yolov5s"
backend: "onnxruntime"
max_batch_size: 0
input [
    {
        name: "images"
        data_type: TYPE_FP32
        dims: [1, 3,640,640]
    }
]
output [
    {
        name: "output0"
        data_type: TYPE_FP32
        dims: [1,25200,85]
    }
]

name：模型的名称，和模型目录名匹配；
backend：指定推理后端，推理后端名称可以参考Backends；
max_batch_size：最大批量输入数目，设置为0表示不限制批量大小；
input/output：指定模型输入和输出大小，每个输入和输出字段包括节点名称、数据类型和数据维度。
- name：节点名称。
- data_type：数据类型，例如TYPE_FP32表示浮点数。
- dims：数据维度，例如[1, 3, 640, 640]表示输入图像的维度。

启动Triton Server

官方推荐通过Docker容器部署NVIDIA Triton Server

# 运行在GPU环境
docker run --gpus=all -it --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v $(pwd):/workdir --workdir=/workdir nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver --model-repository=./model_repository/
# 运行在CPU环境
docker run -it --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v $(pwd):/workdir --workdir=/workdir nvcr.io/nvidia/tritonserver:23.10-py3 tritonserver  --model-repository=./model_repository/

更多的Triton容器版本可以访问NVIDIA GPU Cloud (NGC)，命名格式为nvcr.io/nvidia/tritonserver:<xx.yy>-py3。Triton Server程序安装在容器内的/opt/tritonserver/bin/tritonserver路径下，可以直接在容器内部启动它。

1 2	tritonserver --model-repository=./model_repository/ tritonserver --model-repository=./model_repository/ --http-port=8080 --grpc-port=8090 --metrics-port=8070 # 指定端口号

通过上述命令启动后，Triton Server 会加载模型库中已配置好的模型。如果希望通过 Triton Client 动态加载模型库中的模型，可以在启动时添加 --model-control-mode=explicit参数。这样，Triton Server 在启动后不会自动加载任何模型，而是需要通过 Triton Client 发送模型加载和卸载请求来管理模型。

1	tritonserver --model-repository=./model_repository/ --http-port=8080 --grpc-port=8090 --metrics-port=8070 --model-control-mode=explicit

发送推理请求

安装Triton Client

1	pip install tritonclient[all] -i https://pypi.tuna.tsinghua.edu.cn/simple

CLIENT DEMO

Triton Client的使用流程跟其他推理后端操作类似：首先是创建客户端实例，然后准备输入数据，构建并发送推理请求至Triton Server，最后处理返回的推理结果。

import torch
import cv2
import numpy as np

from tritonclient import grpc as grpcclient  # 导入Triton gRPC客户端模块
from tritonclient.utils import np_to_triton_dtype  # 导入用于转换NumPy数据类型到Triton兼容类型的工具函数

from py.yolov5.yolov5_util import letterbox, non_max_suppression, scale_boxes, draw_results  # 导入YOLOv5相关的预处理和后处理函数
from py.yolov5.general import CLASSES_NAME  # 导入类别名称列表


def preprocess(im0, img_size=640, stride=32, auto=False, device=None, fp16=False):
    # 图像预处理函数，准备输入模型的数据
    im = letterbox(im0, img_size, stride=stride, auto=auto)[0]  # 使用letterbox进行图像填充和调整大小

    im = im.transpose((2, 0, 1))[::-1]  # 将图像从HWC格式转换为CHW格式，并将颜色通道从BGR转换为RGB
    im = np.ascontiguousarray(im)  # 确保数组是连续存储的

    im = torch.from_numpy(im).to(device)  # 将图像数据转换为PyTorch张量并移动到指定设备上
    im = im.half() if fp16 else im.float()  # 如果使用半精度浮点数，则转换张量类型
    im /= 255  # 将像素值从0-255缩放到0.0-1.0
    if len(im.shape) == 3:
        im = im[None]  # 如果是单个样本，则扩展一个维度以表示批次

    return im


def main():
    im0 = cv2.imread("../../assets/bus.jpg")  # 加载测试图片
    im = preprocess(im0, img_size=640, stride=32, auto=False)  # 对图片进行预处理
    im = im.cpu().numpy()  # 将预处理后的图像数据转换回NumPy数组，并移至CPU

    triton_client = grpcclient.InferenceServerClient("localhost:8001")  # 创建Triton推理服务器的gRPC客户端实例

    res = triton_client.is_server_live()  # 检查服务器是否处于活动状态
    res = triton_client.is_server_ready()  # 检查服务器是否准备好接受推理请求

    if not triton_client.is_model_ready(model_name="det_yolov5s"):  # 检查模型是否已经加载完毕
        triton_client.load_model(model_name="det_yolov5s")  # 如果模型未加载，则加载模型
        # triton_client.unload_model(model_name='det_yolov5s')  # 卸载模型（可选）

    metadata = triton_client.get_model_metadata(model_name="det_yolov5s")  # 获取模型元数据
    print(f"metadata:\n{metadata}")  # 打印模型元数据
    model_config = triton_client.get_model_config(model_name='det_yolov5s')  # 获取模型配置
    print(f"model_config:\n{model_config}")  # 打印模型配置

    # res = triton_client.get_model_repository_index()  # 获取模型仓库索引（可选）
    # print(res)  # 打印模型仓库索引（可选）

    # 设置输入
    inputs = [grpcclient.InferInput("images", im.shape, np_to_triton_dtype(np.float32))]  # 定义输入张量及其数据类型
    inputs[0].set_data_from_numpy(im.astype(np.float32))  # 将预处理后的图像数据设置为输入张量的数据

    # 推理
    outputs = [grpcclient.InferRequestedOutput("output0")]  # 定义输出张量
    response = triton_client.infer(model_name="det_yolov5s", inputs=inputs, outputs=outputs)  # 发送推理请求

    # 获取输出
    preds = [response.as_numpy("output0")]  # 从响应中获取输出张量的数据

    pred = non_max_suppression(preds, conf_thres=0.25, iou_thres=0.45, max_det=300)[0]  # 应用非极大值抑制算法过滤预测结果
    print(pred)  # 打印过滤后的预测结果

    boxes = scale_boxes(im.shape[2:], pred[:, :4], im0.shape[:2])  # 将边界框坐标从特征图尺度缩放回原始图像尺度
    confs = pred[:, 4:5]  # 提取置信度分数
    cls_ids = pred[:, 5:6]  # 提取类别ID
    print(f"There are {len(boxes)} objects.")  # 打印检测到的对象数量
    print(boxes)  # 打印边界框坐标

    overlay = draw_results(im0, boxes, confs, cls_ids, CLASSES_NAME, is_xyxy=True)  # 在原始图像上绘制检测结果
    cv2.imshow("overlay", overlay)  # 显示带有检测结果的图像
    cv2.waitKey(0)  # 等待按键事件，保持窗口打开


if __name__ == '__main__':
    main()  # 运行主函数

GRPC vs. HTTP

配置文件

Model Configuration

通过接口get_model_config可以获取到triton处理模型配置文件的全部信息，以yolov5s为例，在Docker CPU环境下启动的配置信息如下：

config {
  name: "det_yolov5s"
  platform: "onnxruntime_onnx"
  version_policy {
    latest {
      num_versions: 1
    }
  }
  input {
    name: "images"
    data_type: TYPE_FP32
    dims: 1
    dims: 3
    dims: 640
    dims: 640
  }
  output {
    name: "output0"
    data_type: TYPE_FP32
    dims: 1
    dims: 25200
    dims: 85
  }
  instance_group {
    name: "det_yolov5s"
    count: 2
    kind: KIND_CPU
  }
  default_model_filename: "model.onnx"
  optimization {
    input_pinned_memory {
      enable: true
    }
    output_pinned_memory {
      enable: true
    }
  }
  backend: "onnxruntime"
}

基本信息
- name: 模型的名称为 "det_yolov5s"。这通常是指模型的唯一标识符。
- platform: 使用的平台是 "onnxruntime_onnx"，这意味着该模型是以 ONNX 格式存储，并使用 ONNX Runtime 进行推理。
版本策略
- version_policy: 定义了模型版本的管理策略。在这个例子中，latest 表示只保留最新的一个版本（num_versions: 1）。
输入配置
- input: 定义了模型的输入。
  - name: 输入张量的名称为 "images"。
  - data_type: 数据类型为 TYPE_FP32，即单精度浮点数。
  - dims: 输入张量的维度为 [1, 3, 640, 640]。这里 1 表示批处理大小，3 表示图像的通道数（RGB），640x640 是图像的尺寸。
输出配置
- output: 定义了模型的输出。
  - name: 输出张量的名称为 "output0"。
  - data_type: 数据类型为 TYPE_FP32。
  - dims: 输出张量的维度为 [1, 25200, 85]。1 同样表示批处理大小，25200 可能是预测框的数量，85 包括了每个预测框的坐标、类别概率等信息。
实例组配置
- instance_group: 定义了模型实例的配置。
  - name: 实例组的名称为 "det_yolov5s"。
  - count: 实例的数量为 2，这意味着 Triton 将会启动两个相同的模型实例来处理请求。
  - kind: 实例运行在 CPU 上 (KIND_CPU)。
默认模型文件名
- default_model_filename: 模型文件的默认名称为 "model.onnx"。
优化选项
- optimization: 定义了一些优化选项以提高推理性能。
  - input_pinned_memory: 启用了输入数据的固定内存 (enable: true)，这样可以加快从主机内存到 GPU 内存的数据传输速度。
  - output_pinned_memory: 同样启用了输出数据的固定内存 (enable: true)。
后端配置
- backend: 指定了使用的后端为 "onnxruntime"，即 ONNX Runtime。

运行在Docker GPU环境下，Triton读取的模型配置信息如下：

config {
  name: "det_yolov5s"
  platform: "onnxruntime_onnx"
  backend: "onnxruntime"
  version_policy {
    latest {
      num_versions: 1
    }
  }
  input {
    name: "images"
    data_type: TYPE_FP32
    dims: 1
    dims: 3
    dims: 640
    dims: 640
  }
  output {
    name: "output0"
    data_type: TYPE_FP32
    dims: 1
    dims: 25200
    dims: 85
  }
  optimization {
    input_pinned_memory {
      enable: true
    }
    output_pinned_memory {
      enable: true
    }
  }
  instance_group {
    name: "det_yolov5s"
    kind: KIND_GPU
    count: 1
    gpus: 0
    gpus: 1
    gpus: 2
    gpus: 3
    gpus: 4
    gpus: 5
    gpus: 6
    gpus: 7
  }
  default_model_filename: "model.onnx"
}

数据类型

Datatypes

TensorRT配置

triton-inference-server/backend

将ONNX模型转换成TensorRT模型，命名为model.plan，存放于文件夹1中。配置文件config.pbtxt内容如下：

name: "det_yolov5s"
platform: "tensorrt_plan"
max_batch_size: 0
input [
    {
        name: "images"
        data_type: TYPE_FP16
        dims: [1, 3,640,640]
    }
]
output [
    {
        name: "output0"
        data_type: TYPE_FP16
        dims: [1,25200,85]
    }
]
version_policy: { latest: { num_versions: 1}}

注：配置文件中的后端配置选项是platform而不是backend

CPU/GPU设置

Instance Groups

使用instance_group可以配置模型在 CPU 或 GPU 环境下的实例数目。以下面的配置为例：

instance_group [
  {
    count: 1
    kind: KIND_GPU
    gpus: [ 0 ]
  },
  {
    count: 2
    kind: KIND_GPU
    gpus: [ 1, 2 ]
  }
]

各个参数解析如下：

count：指定该实例组中的实例数量。
kind：指定实例组的类型，可以是 KIND_CPU 或 KIND_GPU。
gpus：指定实例组使用的 GPU 编号列表。如果 kind 为 KIND_CPU，则不需要指定 gpus。

所以上面的实例组配置中，实现如下功能：

第一个实例组：
- count: 1：创建 1 个实例。
- kind: KIND_GPU：这些实例运行在 GPU 上。
- gpus: [ 0 ]：这些实例使用 GPU 0。
第二个实例组：
- count: 2：创建 2 个实例。
- kind: KIND_GPU：这些实例运行在 GPU 上。
- gpus: [ 1, 2 ]：这些实例分别使用 GPU 1 和 GPU 2。

如果希望在CPU上运行模型，可以配置如下内容：

instance_group [
  {
    count: 3
    kind: KIND_CPU
  }
]

这样模型会在CPU上运行3个实例。

版本策略

Version Policy

在模型仓库中，每个模型指定一个文件夹，该文件夹内除了一个配置文件外，还可以保存多个版本的模型文件，用子文件名表示版本号。格式如下：

注意：版本号用正整数表示，数字越大表示版本越高。

.
└── det_yolov5s
    ├── 1（最小版本号）
    │   └── model.onnx
    ├── 2
    │   └── model.onnx
    ├── 20241119（最高版本号）
    │   └── model.onnx
    └── config.phtxt

NVIDIA Triton在配置文件中使用version_policy控制模型加载版本，同一个模型可以加载多个版本，

方式一：version_policy: { all: {}}，表示加载该模型的所有版本;
方式二：version_policy: { latest: { num_versions: 2}}，表示加载版本最高的前2个模型；
方式三：version_policy: { specific: { versions: [1,3]}}，表示加载指定版本的模型。

默认情况下，版本策略设置为version_policy: { latest: { num_versions: 1}}，即加载版本最高的模型。

大海