Вокруг так много фреймворков для инференса нейронных сетей, что сложно понять, какой именно подойдет тебе лучше всего. Я решил, что реализую одну и ту же задачу на нескольких разных технологиях. Так и родился этот репозиторий. Так как целью проекта являлся инференс, то было решено не обучать что-то с нуля, а взять уже готовые модели.  Хотелось выбрать такую задачу, чтобы она отвечала следующим требованиям: Была не слишком простой (состояла из одной модели), но и не слишком сложной; Содержала нейросети, обученные под разные задачи (например, сетка под детекцию и под сегментацию); Имела реализацию на pytorch (потому что я его знаю). Выбор пал на задачу распознавания российских номеров автомобилей. Модели были взяты из этого репозитория. Пайлпайн распознавания номеров следующий: Детекция номеров с помощью Yolov5; Вырезанные номера прогоняются через Spatial transformer для выравнивания; Текст номера распознается с LPR-net. Останавливаться на том, как работают сетки, не буду (тут есть ссылки на сетки, кому интересно). Так как планируется несколько реализаций инференса, код с ними был вынесен в отдельный пакет, из которого мпортируются сетки. В качестве менеджера зависимостей был выбран Poetry. Хочется, чтобы на данный проект можно было опираться при реализации своего инференса. Для этого нужно, чтобы его можно было запустить и через месяц, и через год. Poetry как раз фиксирует все зависимости (как основные, так и зависимости зависимостей) в специальном файле poetry.lock, что даст воспроизводимость. Веса нейросетей были сохранены прямо в репозитории, потому что: Неохота их куда-то загружать, а потом следить, что ты их случайно не удалил; Так легче воспроизводить (не нужно что-то качать, а потом думать, под каким именем в какую папку положить); Они не так уж много весят. Как их использовать, можно понять из тестов. Для инференса используется Nvidia Triton Inference Server (далее Triton). Я бы сказал, что Triton нацелен на максимальную утилизацию вашего железа: Позволяет передавать данные между моделями без копирования; Не грузит модели, которые не используются; Использует shared memory, чтобы получать данные и отдавать ответы; И другое. Основные недостатки, на мой взгляд: Если запускать на видеокартах, то только на Nvidia; Спроектирован под работу на одной машине. Triton не предусматривает, что вы запустите часть моделей на одной физической машине, а часть на другой; Документация состоит из разбросанных по GitHub-у ридми, часть которых дублируется в сайт. Для использования Triton вам необходимо определить model_repository. Он представляет собой папку с папками под каждую модель. Далее вы сможете объединить модели в некоторый граф вычислений, описав передачу данных между ними. Папка модели должна содержать конфигурацию модели и её веса. Triton поддерживает большое число фреймоворков (бекендов) для обучения сеток. Нам понадобятся Pytorch и Python. Pytorch бекенд позволяет ранить TorchScript модельки. В него у меня переведены LPR-net и STN.  Рассмотрим конвертацию на примере STN. Для этого: 1. Конвертируем в TorchScript 2. Сохраняем в папку с версией модели. Triton хочет, чтобы все ваши модели лежали в отдельной папке с номером версии. Далее вы можете слать запросы на определенную версию. Также с помощью этого механизма можно заменять версии моделек прямо на рабочем проде. 3. Задаем конфигурацию. Основное, что нужно определить - бекенд, батчайсз (чтобы Triton мог процесить батчами), входы и выходы в модель. В нашем случае получаем: Обратите внимание на нижние подчеркивания в названия входов и выходов. Это специальный формат для Pytorch бекенда. По итогу получилась такая папка с моделью Если по какой-то причине конвертация не проходит или есть что-то помимо модели, есть Python бекенд. В нашем случае внутри yolo встроена дополнительная логика, если послать в нее numpy или torch тензор. Чтобы не дописывать дополнительный пре- и постпроцессинг  был использован Python бекенд. Структура папки такая же, как и в предыдущем случае, но вместо весов у нас файл model.py. В нем должен быть реализован класс TritonPythonModel. В нем при инициализации создается модель: Для каждого поступившего в модель запроса получаются предикты и формируются ответы: Чтобы это заработало, необходимо доставить пакетов в интерпретатор, который Triton использует для прогона ваших моделей на бекенде Python. В Docker образах Triton используется Python 3.8. Если у вас другая версия, то нужно собирать под себя специально. Я пошёл проще, и использовал 3.8 🌝. Установить пакеты можно просто через pip при сборке вашего образа. У меня так и не получилось заставить Poetry установить пакеты в глобальный Python, поэтому я немного закостылил и попросил Poetry экспортировать зависимости в requirements.txt. Их я установил через pip. Если у вас не качаются образы тритона с ошибкой 401 unauthorized, то залогиньтесь в nvcr.io Далее также задается конфигурацю. Тут у нас уже два будет два выхода. Батч сайз поставлен 0, так как в препроцессингах для следующих после yolo сеток не реализован случай, если батч больше нуля. Обратите внимание, что если max_batch_size: 0, то вам нужно явно указывать батч димешн. Пайплайн обработки состоит из последовательного применения нескольких моделей. Для задания логики обработки в Triton есть Ensembling и Business Logic Scripting (BLS). Первый позволяет задать логику с помощью конфигов и подходит для простых прямолинейных случаев (например, можно объединить модели STN и LPR-net в одну таким пайплайном), а второй - с помощью кода на Python и подходит для сложных случаев (с циклами, выходами по условию и другому). Было выбрано BLS, так как нужно прекращать обработку, если детектор не нашел номеров, а также делать постобработку выходов LPR-net (распознает текст номера). Для Triton-а наша логика является еще одной моделью с на Python бекенде. В ней мы последовательно передаются данные от модели к модели. Если детектор не нашел номеров, то обработка останавливается. Чтобы передавать данные между моделями, избегая лишних копирований между ГПУ и ЦПУ, используется dlpack. В нашем случае, он позволяет конвертировать torch тензоры в Triton тензоры и обратно. Также задается стандартная конфигурацая для модели на Python бекенде. Triton использует shared memory для Python бекенда и передачи данных в себя. Если вы заводите его в kubernetes, он поинтересуется, куда же она делась. Костыль, который поможет это решить - примонтировать вольюм по пути, где должна быть эта shared memory. Модели заданы. Triton рапортует, что все успешно загружено. Чтобы выполнить описанный пайплайн нужно послать в Triton запрос на использование модели plate_recognition (так названа модель, определяющая наш граф вычислений/пайплайн распознавания номеров). Вы можете написать свои запросы по HTTP или gRPC или воспользоваться готовой библиотекой tritonclient. Примеры можно посмотреть в тестах. Ну вот инференс и написан. Данный пример не слишком простой, чтобы быть в целом бесполезным, но и не слишком сложный, чтобы покрыть все фишки этого фреймоврка для инференса. Советую посмотреть в документации про: Ограничение числа одновременно работающих моделей; Написание моделей, когда не важен порядок выполнения; Кеширования ответов; Анализа скорости работы; И еще много чего. Также в репозиториях Triton-а есть секция examples (например, для Python), в которой есть много полезных примеров. Следующей будет реализация инференса на TorchServe. А пока подписывайтесь на мой телеграм канал. Там рассказываю про нейронки с упором в сервинг.