跳至内容

Triton 推理服务器Ultralytics YOLOv8

Triton Inference Server(原名TensorRT Inference Server)是英伟达公司开发的一款开源软件解决方案。Triton 简化了在生产中大规模部署人工智能模型的过程。将Ultralytics YOLOv8 与Triton Inference Server 集成,可以部署可扩展的高性能深度学习推理工作负载。本指南提供了设置和测试集成的步骤。



观看: NVIDIATriton Inference Server 入门。

什么是Triton Inference Server?

Triton 推理服务器旨在在生产中部署各种人工智能模型。它支持广泛的深度学习和机器学习框架,包括TensorFlow 、PyTorch 、ONNX Runtime 等。它的主要用例包括

  • 从一个服务器实例为多个模型提供服务。
  • 动态加载和卸载模型,无需重启服务器。
  • 集合推理,允许同时使用多个模型来获得结果。
  • 模型版本化,用于 A/B 测试和滚动更新。

先决条件

确保在继续之前具备以下先决条件:

  • 您的计算机上安装了 Docker。
  • 安装 tritonclient:
    pip install tritonclient[all]
    

将YOLOv8 导出为ONNX 格式

在Triton 上部署模型之前,必须将其导出为ONNX 格式。ONNX (Open Neural Network Exchange)是一种允许在不同深度学习框架之间传输模型的格式。使用 export 功能中的 YOLO 类:

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8n.pt")  # load an official model

# Export the model
onnx_file = model.export(format="onnx", dynamic=True)

设置Triton 模型库

Triton 模型库是Triton 可以访问和加载模型的存储位置。

  1. 创建必要的目录结构:

    from pathlib import Path
    
    # Define paths
    model_name = "yolo"
    triton_repo_path = Path("tmp") / "triton_repo"
    triton_model_path = triton_repo_path / model_name
    
    # Create directories
    (triton_model_path / "1").mkdir(parents=True, exist_ok=True)
    
  2. 将导出的ONNX 模型移至Triton 资源库:

    from pathlib import Path
    
    # Move ONNX model to Triton Model path
    Path(onnx_file).rename(triton_model_path / "1" / "model.onnx")
    
    # Create config file
    (triton_model_path / "config.pbtxt").touch()
    

运行Triton 推断服务器

使用 Docker 运行Triton Inference Server:

import contextlib
import subprocess
import time

from tritonclient.http import InferenceServerClient

# Define image https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tritonserver
tag = "nvcr.io/nvidia/tritonserver:23.09-py3"  # 6.4 GB

# Pull the image
subprocess.call(f"docker pull {tag}", shell=True)

# Run the Triton server and capture the container ID
container_id = (
    subprocess.check_output(
        f"docker run -d --rm -v {triton_repo_path}:/models -p 8000:8000 {tag} tritonserver --model-repository=/models",
        shell=True,
    )
    .decode("utf-8")
    .strip()
)

# Wait for the Triton server to start
triton_client = InferenceServerClient(url="localhost:8000", verbose=False, ssl=False)

# Wait until model is ready
for _ in range(10):
    with contextlib.suppress(Exception):
        assert triton_client.is_model_ready(model_name)
        break
    time.sleep(1)

然后使用Triton 服务器模型运行推理:

from ultralytics import YOLO

# Load the Triton Server model
model = YOLO(f"http://localhost:8000/yolo", task="detect")

# Run inference on the server
results = model("path/to/image.jpg")

清理容器:

# Kill and remove the container at the end of the test
subprocess.call(f"docker kill {container_id}", shell=True)

按照上述步骤,您可以在Triton Inference Server 上高效地部署和运行Ultralytics YOLOv8 模型,为深度学习推理任务提供可扩展的高性能解决方案。如果您遇到任何问题或有进一步的疑问,请参阅 Triton 官方文档或向Ultralytics 社区寻求支持。



Created 2023-11-12, Updated 2024-06-02
Authors: glenn-jocher (8)

评论