NVIDIA Triton Inference Server — запуск на Ubuntu

NVIDIA Triton Inference Server представляет собой веб-сервер, на котором запускаются модели распознавания образов. Доступ через http(s) и возможен через GRPC.  Посмотрим, как можно поднять этот сервер, используя docker.

Требования к установленным пакетам на Ubuntu 22.04:

  • NGC Cli
  • Docker 1.5.2
  • nvidia-docker
  • graphics-driver и CUDA.

Клонируем Git репозиторий Triton:

git clone https://github.com/triton-inference-server/server.git

Из клонированного репозитория запускаем команду для загрузки моделей:

~/server/docs/examples$ ./fetch_models.sh

Дальше нужно загрузить нужный докер - смотрим нужную нам версию тут:

https://docs.nvidia.com/deeplearning/triton-inference-server/release-notes/rel-23-12.html#rel-23-12

Запускаем docker:

sudo docker run --gpus=1 --rm -p8000:8000 -p8001:8001 -p8002:8002 -v/home/ubuntu/server/docs/examples/model_repository:/models nvcr.io/nvidia/tritonserver:23.12-py3 tritonserver --model-repository=/models

Должен в конце выдать такой результат:

2024-01-22_16-42-23

Если хотим запустить docker в фоне, то не забываем ключ -d

sudo docker run --gpus=1 -d --rm -p8000:8000 -p8001:8001 -p8002:8002 -v/home/ubuntu/server/docs/examples/model_repository:/models nvcr.io/nvidia/tritonserver:23.12-py3 tritonserver --model-repository=/models

После которой можно посмотреть состояние докеров с помощью:

sudo docker ps -a

2024-01-19_16-56-04

И если надо, то зайти в готовый docker, указав ID контейнера:

sudo docker exec -it e65f17743d60 /bin/bash

Должна появиться командная строка в контейнере тритона:

root@e65f17743d60:/opt/tritonserver#

Но это нужно только для проверки работоспособности и если вы там наберете exit, то docker закроется. Извне докера trition  проверяем его работу:

curl -v localhost:8000/v2/health/ready

2024-01-19_17-03-54

Далее для тестирования загрузим примеры:

sudo docker pull nvcr.io/nvidia/tritonserver:21.12-py3-sdk

И запускаем докер:

sudo docker run -it --rm --net=host nvcr.io/nvidia/tritonserver:21.12-py3-sdk

В нем набираем команду:

/workspace/install/bin/image_client -m densenet_onnx -c 3 -s INCEPTION /workspace/images/mug.jpg

Которая возвращает нам результат классификации:

2024-01-19_17-16-52

 Развертывание Yolo8 

Тритон может использовать модели разных фреймворков, в том числе TensorRT. В примере покажем, как использовать модели ONNX. 

Сначала нужно получить из pt ONNX, что делается просто:

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)

Затем создаем файл config.pbtxt:

name: "yolov8n"

platform: "onnxruntime_onnx"

max_batch_size : 0

input [

  {

    name: "images"

    data_type: TYPE_FP32

    dims: [ 1, 3, 640, 640 ]

  }

]

output [

  {

    name: "output0"

    data_type: TYPE_FP32

    dims: [ -1, -1, -1 ]

  }

]

Если бы использовали TensorRT, то конфигурационный файл был бы таков:

name: "yolov8n"

platform: "tensorrt_plan"

max_batch_size : 0

input [

  {

    name: "images"

    data_type: TYPE_FP16

    dims: [ -1, 3, 640, 640 ]

  }

]

output [

  {

    name: "output0"

    data_type: TYPE_FP16

    dims: [ -1, 84, 8400 ]

  }

]

В model-repository (в нашем примере он находится в папке gitа Tritoна - ~/server/docs/examples) создаем новую папку yolov8n. Внутри нее должно быть так:

yolov8n/

              1/

                            model.onnx

              config.pbtxt

Запускаем докер для тест - модель должна быть готова к использованию:

2024-01-22_16-51-58

Теперь из клиента проверяем запущенный контейнер, для этого нужно установить:

pip install tritonclient[all]

Ну и простенький пример теста модели test.py:

import tritonclient.http as httpclient

import numpy as np

import time

 

client = httpclient.InferenceServerClient(url="0.0.0.0:8000")

inputs = httpclient.InferInput("images", [1,3,640,640], datatype="FP32")

outputs = httpclient.InferRequestedOutput("output0", binary_data=True)

inf_time = 0

for _ in range(10):

    img = np.random.randn(1, 3, 640, 640).astype(np.float32)

   

    inputs.set_data_from_numpy(img, binary_data=True

    # Inference

    start_time = time.time()

    res = client.infer(model_name="yolov8n", inputs=[inputs], outputs=[outputs]).as_numpy('output0')

    end_time = time.time()

    inf_time += (end_time - start_time)

   

print(f"inference time: {inf_time/10 * 1000:.3f} ms")

print(f"input shape: {img.shape}")

print(f"output shape: {res.shape}")

2024-01-22_19-48-31

Инференс работает, но хорошо бы получить изображение и результат. Поэтому усложняем пример:

import tritonclient.http as httpclient

import numpy as np

import time

import cv2

 

colors = np.random.uniform(0, 255, size=(88, 3))

 

def draw_bounding_box(img, class_id, confidence, x, y, x_plus_w, y_plus_h):

    label = str(class_id)

    color = colors[class_id]

    cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)

    cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_DUPLEX, 0.7, color, 1)

    return img

 

 

client = httpclient.InferenceServerClient(url="0.0.0.0:8000")

inputs = httpclient.InferInput("images", [1,3,640,640], datatype="FP32")

outputs = httpclient.InferRequestedOutput("output0", binary_data=True)

inf_time = 0

original_image = cv2.imread("sobaki01.jpg")

print(original_image.shape)

scale_x = original_image.shape[1] / 640

scale_y = original_image.shape[0] / 640

      

resize = cv2.resize(original_image, (640, 640))

img = resize[np.newaxis, :, :, :] / 255.0 

img = img.transpose((0, 3, 1, 2)).astype(np.float32)

   

inputs.set_data_from_numpy(img, binary_data=True

# Inference

start_time = time.time()

res = client.infer(model_name="yolov8n", inputs=[inputs], outputs=[outputs]).as_numpy('output0')

end_time = time.time()

inf_time += (end_time - start_time)

   

print(f"inference time: {inf_time * 1000:.3f} ms")

print(f"input shape: {img.shape}")

print(f"output shape: {res.shape}")

 

# Post Processing

outputs = np.array([cv2.transpose(res[0].astype(np.float32))])

rows = outputs.shape[1]

boxes = [

scores = [

class_ids = [

for i in range(rows):

       classes_scores = outputs[0][i][4:]

       (minScore, maxScore, minClassLoc, (x, maxClassIndex)) = cv2.minMaxLoc(classes_scores)

       if maxScore >= 0.25:

             box = [

                    outputs[0][i][0] - (0.5 * outputs[0][i][2]), outputs[0][i][1] - (0.5 * outputs[0][i][3]),

                    outputs[0][i][2], outputs[0][i][3]]

             boxes.append(box)

             scores.append(maxScore)

             class_ids.append(maxClassIndex)

result_boxes = cv2.dnn.NMSBoxes(boxes, scores, 0.25, 0.45, 0.5)

 

for i in range(len(result_boxes)):

       index = result_boxes[i]

       box = boxes[index]

      

       original_image = draw_bounding_box(original_image, class_ids[index], scores[index], round(box[0] * scale_x), round(box[1] * scale_y),

                                    round((box[0] + box[2]) * scale_x), round((box[1] + box[3]) * scale_y))

                                   

cv2.imwrite("sobaki01out.jpg",original_image)                             

Результат работы примера

На этом все. Также можете посмотреть этот материал по данной теме:

https://blog.kubwa.co.kr/yolov8-with-tensorrt-nvidia-triton-server-f6aa900a53b8