Triton Сервер выводов с Ultralytics YOLOv8

Triton Inference Server (ранее известный как TensorRT Inference Server) - это программное решение с открытым исходным кодом, разработанное компанией NVIDIA. Он представляет собой облачное решение для вычислений, оптимизированное для графических процессоров NVIDIA. Triton упрощает масштабное развертывание моделей ИИ в производстве. Интеграция Ultralytics YOLOv8 с Triton Inference Server позволяет тебе развернуть масштабируемые, высокопроизводительные рабочие нагрузки по анализу глубокого обучения. В этом руководстве описаны шаги по настройке и тестированию интеграции.

Смотри: Начало работы с NVIDIA Triton Inference Server.

Что такое Triton Inference Server?

Triton Inference Server предназначен для развертывания различных моделей ИИ в продакшене. Он поддерживает широкий спектр фреймворков глубокого обучения и машинного обучения, включая 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
   triton_repo_path = Path('tmp') / 'triton_repo'
   triton_model_path = triton_repo_path / 'yolo'
   
   # 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 Inference Server

Запусти сервер Triton Inference Server с помощью Docker:

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 Server:

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)

Выполнив описанные выше шаги, ты сможешь эффективно развернуть и запустить Ultralytics YOLOv8 модели на Triton Inference Server, обеспечив масштабируемое и высокопроизводительное решение для задач глубокого обучения. Если у тебя возникнут какие-либо проблемы или дополнительные вопросы, обратись к официальной документации Triton или обратись за поддержкой к сообществу Ultralytics .

Комментарии