部署 ONNX 模型#

本 README 展示了如何在 Triton 推理服务器上部署一个简单的 ResNet 模型。

步骤 1:设置 Triton 推理服务器#

要使用 Triton,我们需要构建一个模型仓库。仓库的结构如下

model_repository
|
+-- resnet
    |
    +-- config.pbtxt
    +-- 1
        |
        +-- model.onnx

config.pbtxt 配置文件是可选的。如果用户未提供配置文件,则 Triton 推理服务器会自动生成配置文件。如果您是 Triton 新手,强烈建议查看概念指南的第 1 部分

mkdir -p model_repository/densenet_onnx/1
wget -O model_repository/densenet_onnx/1/model.onnx \
     https://contentmamluswest001.blob.core.windows.net/content/14b2744cf8d6418c87ffddc3f3127242/9502630827244d60a1214f250e3bbca7/08aed7327d694b8dbaee2c97b8d0fcba/densenet121-1.2.onnx

docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${PWD}/model_repository:/models nvcr.io/nvidia/tritonserver:<yy.mm>-py3 tritonserver --model-repository=/models

步骤 3:使用 Triton 客户端查询服务器#

安装依赖项并下载示例图像以测试推理。

docker run -it --rm --net=host -v ${PWD}:/workspace/ nvcr.io/nvidia/tritonserver:<yy.mm>-py3-sdk bash
pip install torchvision

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

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

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

其次,我们指定模型的输入和输出层名称,并描述预期输入的形状和数据类型。

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

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

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

# Querying the server
results = client.infer(model_name="densenet_onnx", inputs=[inputs], outputs=[outputs])
inference_output = results.as_numpy('fc6_1').astype(str)

print(np.squeeze(inference_output)[:5])

相同的输出应如下所示

['11.549026:92' '11.232335:14' '7.528014:95' '6.923391:17' '6.576575:88']

此处的输出格式为 <confidence_score>:<classification_index>。要了解如何将这些映射到标签名称以及更多信息,请参阅我们的文档。上面的客户端代码在 client.py 中提供。