使用 Triton 部署 Torch-TensorRT 模型¶
在討論機器學習基礎架構時，最佳化和部署是密不可分的。一旦完成網路層級的最佳化以獲得最大效能，下一步就是部署它。
然而，部署這個最佳化模型也有它自己的一系列考量和挑戰，例如：建立支援並行模型執行的基礎架構、支援透過 HTTP 或 gRPC 的客戶端等等。
Triton 推論伺服器解決了上述問題以及更多。讓我們一步一步地討論使用 Torch-TensorRT 最佳化模型、將其部署在 Triton 推論伺服器上以及建立客戶端以查詢模型的過程。
步驟 1：使用 Torch-TensorRT 最佳化您的模型¶
大多數 Torch-TensorRT 使用者都會熟悉這個步驟。為了示範起見，我們將使用來自 Torchhub 的 ResNet50 模型。
我們將在 //examples/triton 目錄中工作，該目錄包含本教學中使用的腳本。
首先，拉取 NGC PyTorch Docker 容器。您可能需要建立一個帳戶並從 這裡取得 API 金鑰。註冊並使用您的金鑰登入（註冊後按照這裡的說明）。
# YY.MM is the yy:mm for the publishing tag for NVIDIA's Pytorch
# container; eg. 24.08
# NOTE: Use the publishing tag for both the PyTorch container and the Triton Containers
docker run -it --gpus all -v ${PWD}:/scratch_space nvcr.io/nvidia/pytorch:YY.MM-py3
cd /scratch_space
有了容器，我們可以將模型匯出到我們 Triton 模型儲存庫中的正確目錄。這個匯出腳本使用 Torch-TensorRT 的 Dynamo 前端將 PyTorch 模型編譯為 TensorRT。然後，我們使用 TorchScript 作為序列化格式來儲存模型，Triton 支援這種格式。
import torch
import torch_tensorrt as torchtrt
import torchvision
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_tensorrt.dtype.f16}
)
ts_trt_model = torch.jit.trace(trt_model, torch.rand(1, 3, 224, 224).to("cuda"))
# Save the model
torch.jit.save(ts_trt_model, "/triton_example/model_repository/resnet50/1/model.pt")
您可以使用以下命令執行腳本（從 //examples/triton）：
docker run --gpus all -it --rm -v ${PWD}:/triton_example nvcr.io/nvidia/pytorch:YY.MM-py3 python /triton_example/export.py
這會將 ResNet 模型的序列化 TorchScript 版本儲存在模型儲存庫的正確目錄中。
步驟 2：設定 Triton 推論伺服器¶
如果您是 Triton 推論伺服器的新手，並且想了解更多資訊，我們強烈建議您查看我們的 Github 儲存庫。
要使用 Triton，我們需要建立一個模型儲存庫。顧名思義，模型儲存庫是推論伺服器託管的模型的儲存庫。雖然 Triton 可以從多個儲存庫提供模型，但在本範例中，我們將討論模型儲存庫的最簡單形式。
此儲存庫的結構應如下所示
model_repository
|
+-- resnet50
|
+-- config.pbtxt
+-- 1
|
+-- model.pt
Triton 需要兩個檔案才能提供模型：模型本身和模型設定檔，該檔案通常以 config.pbtxt 形式提供。對於我們在步驟 1 中準備的模型，可以使用以下配置
name: "resnet50"
backend: "pytorch"
max_batch_size : 0
input [
{
name: "x"
data_type: TYPE_FP32
dims: [ 1, 3, 224, 224 ]
}
]
output [
{
name: "output0"
data_type: TYPE_FP32
dims: [1, 1000]
}
]
config.pbtxt 檔案用於描述確切的模型配置，其中包含輸入和輸出層的名稱和形狀、資料類型、排程和批次處理詳細資訊等詳細資訊。如果您是 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. Roughly, like publishing tags should have the same TensorRT version
docker run --gpus all --rm -p 8000:8000 -p 8001:8001 -p 8002:8002 -v ${PWD}:/triton_example nvcr.io/nvidia/tritonserver:YY.MM-py3 tritonserver --model-repository=/triton_example/model_repository
這應該會啟動 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="/triton_example/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).unsqueeze(0).numpy()
transformed_img = rn50_preprocess()
建立客戶端需要三個基本點。首先，我們建立與 Triton 推論伺服器的連線。
# Setting up client
client = httpclient.InferenceServerClient(url="localhost:8000")
其次，我們指定模型的輸入和輸出層的名稱。這可以在匯出期間取得，並且應該已在您的 config.pbtxt 中指定
inputs = httpclient.InferInput("x", transformed_img.shape, datatype="FP32")
inputs.set_data_from_numpy(transformed_img, binary_data=True)
outputs = httpclient.InferRequestedOutput("output0", 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('output0')
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 推論伺服器的文件。
您可以使用以下命令快速嘗試此客戶端
# Remember to use the same publishing tag for all steps (e.g. 24.08)
docker run -it --net=host -v ${PWD}:/triton_example nvcr.io/nvidia/tritonserver:YY.MM-py3-sdk bash -c "pip install torchvision && python /triton_example/client.py"
