PyTorch/TensorRT模型在Triton推理服务器上的部署指南

PyTorch/TensorRT模型在Triton推理服务器上的部署指南TensorRTPyTorch/TorchScript/FX compiler for NVIDIA GPUs using TensorRT项目地址: ...

魏秦任

406人浏览 · 2025-06-20 09:49:32

魏秦任 · 2025-06-20 09:49:32 发布

PyTorch/TensorRT模型在Triton推理服务器上的部署指南

TensorRT PyTorch/TorchScript/FX compiler for NVIDIA GPUs using TensorRT 项目地址: https://gitcode.com/gh_mirrors/te/TensorRT

前言

在机器学习生产环境中，模型优化与部署是两个紧密相连的关键环节。PyTorch/TensorRT项目通过将PyTorch模型转换为TensorRT引擎，显著提升了模型推理性能。而如何高效地部署这些优化后的模型，则需要借助专业的推理服务器。本文将详细介绍如何使用Triton推理服务器来部署经过Torch-TensorRT优化的模型。

环境准备

在开始之前，我们需要准备以下环境：

支持CUDA的NVIDIA显卡
Docker环境
NVIDIA容器工具包

建议使用NGC提供的PyTorch容器，它已经预装了PyTorch、TensorRT等必要的软件栈。

第一步：使用Torch-TensorRT优化模型

Torch-TensorRT是PyTorch的一个扩展，它能够将PyTorch模型转换为TensorRT引擎，同时保持PyTorch的易用性。我们以ResNet50为例，演示如何优化模型：

import torch
import torch_tensorrt
import torchvision

# 加载预训练的ResNet50模型
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True).eval().to("cuda")

# 使用Torch-TensorRT进行编译
trt_model = torch_tensorrt.compile(
    model,
    inputs=[torch_tensorrt.Input((1, 3, 224, 224))],
    enabled_precisions={torch_tensorrt.dtype.f16}  # 使用FP16精度
)

# 转换为TorchScript格式
ts_trt_model = torch.jit.trace(trt_model, torch.rand(1, 3, 224, 224).to("cuda"))

# 保存模型
torch.jit.save(ts_trt_model, "model_repository/resnet50/1/model.pt")

这段代码完成了几个关键步骤：

加载预训练的ResNet50模型
使用Torch-TensorRT进行编译优化，指定输入尺寸和计算精度
将优化后的模型转换为TorchScript格式
保存模型到Triton模型仓库的指定位置

第二步：配置Triton推理服务器

Triton推理服务器需要一个结构化的模型仓库来管理模型。对于我们的ResNet50模型，仓库结构如下：

model_repository/
└── resnet50/
    ├── config.pbtxt
    └── 1/
        └── model.pt

关键的配置文件config.pbtxt内容如下：

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]
  }
]

配置文件中的关键参数说明：

backend: 指定使用PyTorch后端
max_batch_size: 设为0表示禁用动态批处理
input/output: 定义模型的输入输出张量信息

启动Triton服务器的命令如下：

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

第三步：构建客户端应用程序

客户端需要完成以下工作：

图像预处理
与Triton服务器建立连接
发送推理请求并处理结果

以下是完整的客户端代码示例：

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

def preprocess_image(img_path):
    """图像预处理函数"""
    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()

# 1. 预处理输入图像
input_data = preprocess_image("img1.jpg")

# 2. 创建Triton客户端
client = httpclient.InferenceServerClient(url="localhost:8000")

# 3. 准备输入输出
inputs = httpclient.InferInput("x", input_data.shape, "FP32")
inputs.set_data_from_numpy(input_data)

outputs = httpclient.InferRequestedOutput("output0")

# 4. 发送推理请求
response = client.infer("resnet50", inputs=[inputs], outputs=[outputs])

# 5. 处理结果
results = response.as_numpy("output0")
top5_indices = np.argsort(results[0])[-5:][::-1]
print("Top 5 predictions:", top5_indices)