Вокруг так много фреймворков для инференса нейронок, что глаза разбегаются. Продолжаем цикл о реализации сервинга одной задачи, но разными инструментами. В прошлый раз реализация была на 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.
Возможности:
Поддержка нескольких форматов моделей (torchscript, onnx, ipex, tensorrt);
Объединение нескольких моделей в один граф/workflow;
Инференс API (REST и GRPC);
API для управления моделями;
Метрики из коробки.
Конвертируем модели
Как и Triton, TorchServe требует от пользователя перевести модели в свой формат. Для этого есть утилиты torch-model-archiver и torch-workflow-archiver для моделей и графов соответственно.
Для конвертации нам нужно:
Модель в формате TorchServe/Onnx/др.;
Скрипт, описывающий пайплайн работы модели.
Такой скрипт называется 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, поэтому советую посмотреть в документации про:
Получение explanations;
Снятие метрик работающего сервера;
Подписывайтесь на мой канал - там я рассказываю про нейронки с упором в сервинг.