Вокруг так много фреймворков для инференса нейронок, что глаза разбегаются. Продолжаем ��икл о реализации сервинга одной задачи, но разными инструментами. В прошлый раз реализация была на Nvidia Triton Inference Serve (за анонсами прошу в мой телеграм канал. Код к статье находится в репозитории.

Задача

В качестве задачи взято распознавание российских автомобильных номеров. Модели были взяты из этого репозитория.

Пайплайн распознавания следующий:
1. Детекция номеров с помощью Yolov5;
2. Вырезанные номера прогоняются через Spatial transformer (STN) для выравнивания;
3. Текст номера распознается с LPR-net.

Фреймворк

Для инференса используется [TorchServe]. Данный фреймворк является частью экосистемы Pytorch. Он активно развивается.

В документации о нем говорится следующее:

TorchServe is a performant, flexible and easy to use tool for serving PyTorch eager mode and torschripted models.

Возможности:

Конвертируем модели

Как и Triton, TorchServe требует от пользователя перевести модели в свой формат. Для этого есть утилиты torch-model-archiver и torch-workflow-archiver для моделей и графов соответственно.

Для конвертации нам нужно:

  1. Модель в формате TorchServe/Onnx/др.;

  2. Скрипт, описывающий пайплайн работы модели.

Такой скрипт называется handler. В нем определяются основные этапы жизненного цикла модели (инициализация, предобработка, предсказание, постобработка и др.). Для типовых задач они уже предопределены.

Модели STN и LPR легко конверуются в TorchServe, поэтому в их хэндлерах не используются дополнительные библиотеки. Импорты выглядят так:

import json
import logging
from abc import ABC
import numpy as np
import torch
from ts.torch_handler.base_handler import BaseHandler

Yolo нельзя было просто перевести в TorchScript, так как часть логики для обработки запросов оставалась снаружи модели. Так как копаться с этим желания не было, а также ради более приближенного к жизни сценария, в хэндлере модели Yolo инициализируется из TorchHub. В импортах мы уже видим и сторонние модули:

from inference_torchserve.data_models import PlatePrediction
from nn.inference.predictor import prepare_detection_input, prepare_recognition_input
from nn.models.yolo import load_yolo
from nn.settings import settings

Чтобы это работало, необходимо в докерфайле установить в глобальный интерпретатор необходимые вам пакеты.

В TorchServe не нужно жестко задавать тип и размерность входов и выходов модели, поэтому никакие конфиги для моделей определять не нужно. С одной стороны это удобно, а с другой - порождает хаос, если не следовать какому-то одному формату.

Конвертированная модель представляет собой zip архив с расширением .mar, в котором лежат все артефакты (служебная информация, веса, скрипты и дополнительные файлы).

.
├── MAR-INF
│   └── MANIFEST.json
├── stn.pt
└── stn.py

На мой взгляд, решение с архивом неудобно для разработки. После любого изменения необходимо заново конвертировать модель. Также я испытывал проблемы при запуске в нем удаленного дебагера.

Чтобы TorchServe загрузил модели, их нужно положить в одну папку - model storage и указать путь до нее в параметрах. Чтобы при запуске поднимались все модели, необходимо указать --models all.

Делаем пайплайн распознавания

Выбранный пайплайн распознавания номера автомобиля состоит из последовательного предсказания несколькими моделями. Для этого в TorchServe есть Workflow. Он позв��ляет задать как последовательный, так и параллельный граф обработки:

# последовательный 
dag:
  pre_processing : [m1]
  m1 : [m2]
  m2 : [postprocessing]

input -> function1 -> model1 -> model2 -> function2 -> output

# параллельный граф
dag:
  pre_processing: [model1, model2]
  model1: [aggregate_func]
  model2: [aggregate_func]

                          model1
                         /       \
input -> preprocessing ->         -> aggregate_func
                         \       /
                          model2

Для рассматриваемой задачи получился следующий последовательно-параллельный граф. Узел aggregate объединяет координаты номеров с распознанными текстами.

    ┌──────┐
    │ YOLO ├─────┐
    └──┬───┘     │
       │         v
       │      ┌─────┐
plate  │      │ STN │
coords │      └──┬──┘
       │         │
       │         v
       │      ┌──────┐
       │      │LPRNET│
       │      └──┬───┘
       v         │
   ┌─────────┐   │ plate
   │aggregate│<──┘ texts
   └─────────┘

Для удобства и простоты данные между моделями передаются в виде словарей. Сериализация таких данных в TorchServe весьма неэффективна (переводят в строку и добавляют переносов строк), поэтому старайтесь передавать их как тензоры или байты.

Учтите, что workflow нельзя стартовать автоматически при запуске сервера - необходимо явно послать запрос на это. Если очень хочется делать при поднятии сервера, то можно так.

curl -X POST http://localhost:8081/workflows?url=plate_recognition

Использование моделей

Модели определены. Сервер запущен.

Чтобы выполнить определенную ранее модель или workflow нужно послать в TorchServe запрос на использование plate_recognition (я пользовался REST, но есть еще и GRPC). Для моделей используется эндпоинт predictions, а для workflow wfpredict.

response = requests.post(
    "http://localhost:8080/predictions/yolo", data=image.open("rb").read()
)

response = requests.post(
    "http://localhost:8080/wfpredict/plate_recognition", 
    data=image.open("rb").read()
)

Заключение

Ну вот инференс и написан. Данный пример не слишком простой, чтобы быть в целом бесполезным, но и не слишком сложный, чтобы покрыть все фишки этого фреймоврка для инференса.

В этом туториале были раскрыты не все возможности TorchServe, поэтому советую посмотреть в документации про:

Подписывайтесь на мой канал - там я рассказываю про нейронки с упором в сервинг.