使用 Triton 部署 Torch-TensorRT 模型

在讨论机器学习基础设施时，优化和部署是密不可分的。完成网络级别的优化以获得最大性能后，下一步就是进行部署。

但是，部署此优化后的模型会带来自身的一系列考虑因素和挑战，例如：构建基础设施以支持并发模型执行、通过 HTTP 或 gRPC 支持客户端等等。

Triton 推理服务器 解决上述问题以及更多问题。让我们逐步讨论使用 Torch-TensorRT 优化模型、将其部署在 Triton 推理服务器上以及构建客户端以查询模型的过程。

步骤 1：使用 Torch-TensorRT 优化您的模型

大多数 Torch-TensorRT 用户都熟悉此步骤。为了演示的目的，我们将使用来自 Torchhub 的 ResNet50 模型。

首先，让我们拉取 NGC PyTorch Docker 容器。您可能需要创建一个帐户并从 这里 获取 API 密钥。注册并使用您的密钥登录（注册后按照 此处 的说明进行操作）。

# <xx.xx> is the yy:mm for the publishing tag for NVIDIA's Pytorch
# container; eg. 22.04

docker run -it --gpus all -v ${PWD}:/scratch_space nvcr.io/nvidia/pytorch:<xx.xx>-py3
cd /scratch_space

进入容器后，我们可以继续从 Torchhub 下载 ResNet 模型并使用 Torch-TensorRT 进行优化。

import torch
import torch_tensorrt
torch.hub._validate_not_a_forked_repo=lambda a,b,c: True

# load model
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True).eval().to("cuda")

# Compile with Torch TensorRT;
trt_model = torch_tensorrt.compile(model,
    inputs= [torch_tensorrt.Input((1, 3, 224, 224))],
    enabled_precisions= { torch.half} # Run with FP32
)

# Save the model
torch.jit.save(trt_model, "model.pt")

复制模型后，退出容器。流程中的下一步是设置 Triton 推理服务器。

步骤 2：设置 Triton 推理服务器

如果您是 Triton 推理服务器的新手，想要了解更多信息，我们强烈建议您查看我们的 Github 仓库。

要使用 Triton，我们需要创建一个模型仓库。顾名思义，模型仓库是推理服务器托管的模型的仓库。虽然 Triton 可以从多个仓库提供服务，但在本示例中，我们将讨论模型仓库的最简单形式。

此仓库的结构应如下所示

model_repository
|
+-- resnet50
    |
    +-- config.pbtxt
    +-- 1
        |
        +-- model.pt

Triton 需要两个文件来提供服务模型：模型本身和模型配置文件，该文件通常以 config.pbtxt 格式提供。对于我们在步骤 1 中准备的模型，可以使用以下配置

name: "resnet50"
platform: "pytorch_libtorch"
max_batch_size : 0
input [
  {
    name: "input__0"
    data_type: TYPE_FP32
    dims: [ 3, 224, 224 ]
    reshape { shape: [ 1, 3, 224, 224 ] }
  }
]
output [
  {
    name: "output__0"
    data_type: TYPE_FP32
    dims: [ 1, 1000 ,1, 1]
    reshape { shape: [ 1, 1000 ] }
  }
]

config.pbtxt 文件用于描述精确的模型配置，其中包含输入和输出层（s）的名称和形状、数据类型、调度和批处理详细信息等等。如果您是 Triton 的新手，我们强烈建议您查看我们的 文档部分 以获取更多详细信息。

完成模型仓库的设置后，我们可以继续使用以下 docker 命令启动 Triton 服务器。有关容器的拉取标签，请参阅 此页面。

# Make sure that the TensorRT version in the Triton container
# and TensorRT version in the environment used to optimize the model
# are the same.

docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v /full/path/to/the_model_repository/model_repository:/models nvcr.io/nvidia/tritonserver:<xx.yy>-py3 tritonserver --model-repository=/models

这应该会启动 Triton 推理服务器。下一步，构建一个简单的 http 客户端来查询服务器。

步骤 3：构建 Triton 客户端以查询服务器

在继续之前，请确保您手边有一个示例图像。如果您没有，请下载一个示例图像以测试推理。在本节中，我们将介绍一个非常基本的客户端。有关各种更完善的示例，请参阅 Triton 客户端仓库

wget  -O img1.jpg "https://www.hakaimagazine.com/wp-content/uploads/header-gulf-birds.jpg"

然后，我们需要安装用于构建 Python 客户端的依赖项。这些将因客户端而异。有关 Triton 支持的所有语言的完整列表，请参阅 Triton 的客户端仓库。

pip install torchvision
pip install attrdict
pip install nvidia-pyindex
pip install tritonclient[all]

让我们进入客户端。首先，我们编写一个小的预处理函数来调整查询图像的大小并进行归一化。

import numpy as np
from torchvision import transforms
from PIL import Image
import tritonclient.http as httpclient
from tritonclient.utils import triton_to_np_dtype

# preprocessing function
def rn50_preprocess(img_path="img1.jpg"):
    img = Image.open(img_path)
    preprocess = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return preprocess(img).numpy()

transformed_img = rn50_preprocess()

构建客户端需要三个基本点。首先，我们与 Triton 推理服务器建立连接。

# Setting up client
client = httpclient.InferenceServerClient(url="localhost:8000")

其次，我们指定模型的输入和输出层（s）的名称。

inputs = httpclient.InferInput("input__0", transformed_img.shape, datatype="FP32")
inputs.set_data_from_numpy(transformed_img, binary_data=True)

outputs = httpclient.InferRequestedOutput("output__0", binary_data=True, class_count=1000)

最后，我们向 Triton 推理服务器发送推理请求。

# Querying the server
results = client.infer(model_name="resnet50", inputs=[inputs], outputs=[outputs])
inference_output = results.as_numpy('output__0')
print(inference_output[:5])

输出应如下所示

[b'12.468750:90' b'11.523438:92' b'9.664062:14' b'8.429688:136'
 b'8.234375:11']

此处的输出格式为 <confidence_score>:<classification_index>。要了解如何将其映射到标签名称等等，请参阅 Triton 推理服务器的 文档。