Делаем инференс на Nvidia Triton Inference Server

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

Задача

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

Хотелось выбрать такую задачу, чтобы она отвечала следующим требованиям:

• Была не слишком простой (состояла из одной модели), но и не слишком сложной;

• Содержала нейросети, обученные под разные задачи (например, сетка под детекцию и под сегментацию);

• Имела реализацию на pytorch (потому что я его знаю).

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

Пайлпайн распознавания номеров следующий:

1. Детекция номеров с помощью Yolov5;

2. Вырезанные номера прогоняются через Spatial transformer для выравнивания;

3. Текст номера распознается с LPR-net.

Останавливаться на том, как работают сетки, не буду (тут есть ссылки на сетки, кому интересно).

Нейросети

Так как планируется несколько реализаций инференса, код с ними был вынесен в отдельный пакет, из которого мпортируются сетки.

В качестве менеджера зависимостей был выбран Poetry. Хочется, чтобы на данный проект можно было опираться при реализации своего инференса. Для этого нужно, чтобы его можно было запустить и через месяц, и через год. Poetry как раз фиксирует все зависимости (как основные, так и зависимости зависимостей) в специальном файле poetry.lock, что даст воспроизводимость.

Веса нейросетей были сохранены прямо в репозитории, потому что:

1. Неохота их куда-то загружать, а потом следить, что ты их случайно не удалил;

2. Так легче воспроизводить (не нужно что-то качать, а потом думать, под каким именем в какую папку положить);

3. Они не так уж много весят.

Как их использовать, можно понять из тестов.

Фреймворк

Для инференса используется Nvidia Triton Inference Server (далее Triton).

Я бы сказал, что Triton нацелен на максимальную утилизацию вашего железа:

• Позволяет передавать данные между моделями без копирования;

• Не грузит модели, которые не используются;

• Использует shared memory, чтобы получать данные и отдавать ответы;

• И другое.

Основные недостатки, на мой взгляд:

• Работает только на видеокартах Nvidia;

• Спроектирован под работу на одной машине. Triton не предусматривает, что вы запустите часть моделей на одной физической машине, а часть на другой;

• Документация состоит из разбросанных по GitHub-у ридми, часть которых дублируется в сайт.

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

Pytorch модели

Для использования Triton вам необходимо определить model_repository. Он представляет собой папку с папками под каждую модель. Далее вы сможете объединить модели в некоторый граф вычислений, описав передачу данных между ними.

Папка модели должна содержать конфигурацию модели и её веса.

Triton поддерживает большое число фреймоворков (бекендов) для обучения сеток. Нам понадобятся Pytorch и Python.

Pytorch бекенд позволяет ранить TorchScript модельки. В него у меня переведены LPR-net и STN.

Рассмотрим конвертацию на примере STN. Для этого:

1. Конвертируем в TorchScript

dummy_input = torch.randn(1, 3, 24, 94, device=device) model = load_stn(settings.STN.WEIGHTS, device) traced_model = torch.jit.trace(model, [dummy_input])

2. Сохраняем в папку с версией модели. Triton хочет, чтобы все ваши модели лежали в отдельной папке с номером версии. Далее вы можете слать запросы на определенную версию. Также с помощью этого механизма можно заменять версии моделек прямо на рабочем проде.

model_path = (         Path(__file__).resolve().parents[1]         / "model_repository"         / "stn"         / "1"         / "model.pt"     ) traced_model.save(model_path)

3. Задаем конфигурацию. Основное, что нужно определить — бекенд, батчайсз (чтобы Triton мог процесить батчами), входы и выходы в модель. В нашем случае получаем:

backend: "pytorch" max_batch_size: 32 input [     {         name: "input__0"         data_type: TYPE_FP32         dims: [ 3, 24, 94 ]     } ] output [     {         name: "output__0"         data_type: TYPE_FP32         dims: [ 3, 24, 94 ]     } ]

Hidden text

Обратите внимание на нижние подчеркивания в названия входов и выходов. Это специальный формат для Pytorch бекенда.

По итогу получилась такая папка с моделью

Python модели

Если по какой-то причине конвертация не проходит или есть что-то помимо модели, есть Python бекенд. В нашем случае внутри yolo встроена дополнительная логика, если послать в нее numpy или torch тензор. Чтобы не дописывать дополнительный пре- и постпроцессинг был использован Python бекенд.

Структура папки такая же, как и в предыдущем случае, но вместо весов у нас файл model.py. В нем должен быть реализован класс TritonPythonModel. В нем при инициализации создается модель:

def initialize(self, args):     self.model_config = json.loads(args["model_config"])     self.model = load_yolo(         settings.YOLO.WEIGHTS, settings.YOLO.CONFIDENCE, torch.device("cuda")     )

Для каждого поступившего в модель запроса получаются предикты и формируются ответы:

def execute(     self, requests: List[pb_utils.InferenceRequest] ) -> List[pb_utils.InferenceRequest]:      responses = []      for request in requests:         image = pb_utils.get_input_tensor_by_name(request, "input__0").as_numpy()         image = prepare_detection_input(image)          detection = self.model(image, size=settings.YOLO.PREDICT_SIZE)          df_results = detection.pandas().xyxy[0]         img_plates = prepare_recognition_input(             df_results, image, return_torch=False         )          out_tensor_0 = pb_utils.Tensor(             "output__0", img_plates.astype(self.output0_dtype)         )         out_tensor_1 = pb_utils.Tensor(             "output__1",             df_results[["xmin", "ymin", "xmax", "ymax"]]             .to_numpy()             .astype(self.output1_dtype),         )          inference_response = pb_utils.InferenceResponse(             output_tensors=[out_tensor_0, out_tensor_1]         )         responses.append(inference_response)      return responses

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

В Docker образах Triton используется Python 3.8. Если у вас другая версия, то нужно собирать под себя специально. Я пошёл проще, и использовал 3.8 ?.

Установить пакеты можно просто через pip при сборке вашего образа. У меня так и не получилось заставить Poetry установить пакеты в глобальный Python, поэтому я немного закостылил и попросил Poetry экспортировать зависимости в requirements.txt. Их я установил через pip.

Hidden text

Если у вас не качаются образы тритона с ошибкой 401 unauthorized, то залогиньтесь в nvcr.io

Далее также задается конфигурацю. Тут у нас уже два будет два выхода. Батч сайз поставлен 0, так как в препроцессингах для следующих после yolo сеток не реализован случай, если батч больше нуля.

Hidden text

Обратите внимание, что если max_batch_size: 0, то вам нужно явно указывать батч димешн.

backend: "python" max_batch_size: 0 input [     {         name: "input__0"         data_type: TYPE_UINT8         dims: [ -1, -1, 3 ]     } ] output [     {         name: "output__0"         data_type: TYPE_FP32         dims: [ -1, 3, 24, 94 ]     },     {         name: "output__1"         data_type: TYPE_FP32         dims: [ -1, 4 ]     } ]

Граф вычислений

Пайплайн обработки состоит из последовательного применения нескольких моделей. Для задания логики обработки в Triton есть Ensembling и Business Logic Scripting (BLS).

Первый позволяет задать логику с помощью конфигов и подходит для простых прямолинейных случаев (например, можно объединить модели STN и LPR-net в одну таким пайплайном), а второй — с помощью кода на Python и подходит для сложных случаев (с циклами, выходами по условию и другому).

Было выбрано BLS, так как нужно прекращать обработку, если детектор не нашел номеров, а также делать постобработку выходов LPR-net (распознает текст номера).

Для Triton-а наша логика является еще одной моделью с на Python бекенде.

В ней мы последовательно передаются данные от модели к модели. Если детектор не нашел номеров, то обработка останавливается.

cropped_images, coordinates = self.predict(     model_name="yolo",     ... )  num_plates = from_dlpack(coordinates.to_dlpack()).shape[0] if num_plates == 0:     ...     continue  (plate_features,) = self.predict(     model_name="stn",     ... )  (text_features,) = self.predict(     model_name="lprnet",     ... )

Hidden text

Чтобы передавать данные между моделями, избегая лишних копирований между ГПУ и ЦПУ, используется dlpack. В нашем случае, он позволяет конвертировать torch тензоры в Triton тензоры и обратно.

Также задается стандартная конфигурацая для модели на Python бекенде.

backend: "python" max_batch_size: 0 input [     {         name: "input__0"         data_type: TYPE_UINT8         dims: [ -1, -1, 3 ]     } ] output [     {         name: "coordinates"         data_type: TYPE_FP32         dims: [ -1, 4 ]     },     {         name: "texts"         data_type: TYPE_STRING         dims: [ -1, 1 ]     } ]

Hidden text

Triton использует shared memory для Python бекенда и передачи данных в себя. Если вы заводите его в kubernetes, он поинтересуется, куда же она делась. Костыль, который поможет это решить — примонтировать вольюм по пути, где должна быть эта shared memory.

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

Модели заданы. Triton рапортует, что все успешно загружено.

Чтобы выполнить описанный пайплайн нужно послать в Triton запрос на использование модели plate_recognition (так названа модель, определяющая наш граф вычислений/пайплайн распознавания номеров). Вы можете написать свои запросы по HTTP или gRPC или воспользоваться готовой библиотекой tritonclient. Примеры можно посмотреть в тестах.

triton_client = httpclient.InferenceServerClient(url="localhost:8000") model_name = "plate_recognition" image = cv2.imread(img_path)  inputs, outputs = [], []  inputs.append(httpclient.InferInput("input__0", image.shape, "UINT8")) inputs[0].set_data_from_numpy(image)  outputs.append(httpclient.InferRequestedOutput("coordinates", binary_data=False)) outputs.append(httpclient.InferRequestedOutput("texts", binary_data=False))  results = triton_client.infer(     model_name=model_name,     inputs=inputs,     outputs=outputs, )  coordinates = results.as_numpy("coordinates") texts = results.as_numpy("texts")

Заключение

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

Советую посмотреть в документации про:

• Ограничение числа одновременно работающих моделей;

• Написание моделей, когда не важен порядок выполнения;

• Кеширования ответов;

• Анализа скорости работы;

• И еще много чего.

Также в репозиториях Triton-а есть секция examples (например, для Python), в которой есть много полезных примеров.

Следующей будет реализация инференса на TorchServe. А пока подписывайтесь на мой телеграм канал. Там рассказываю про нейронки с упором в сервинг.