部署 PyTorch 模型#

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

步骤 1：导出模型#

保存 PyTorch 模型。此模型需要被追踪/脚本化以获得 torchscript 模型。

# <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}:/workspace nvcr.io/nvidia/pytorch:<xx.xx>-py3
python export.py

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

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

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

此演示包含一个模型配置示例，文件名为 config.pbtxt。如果您是 Triton 新手，强烈建议查看概念指南的第 1 部分。

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

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

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

docker run -it --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("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])
predictions = results.as_numpy('output__0')
print(predictions[: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>。要了解如何将这些映射到标签名称以及更多信息，请参阅我们的文档。上面的客户端代码在 client.py 中提供。

在本页中