Делаем инференс на 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 ]
    }
]

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

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.

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

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",
    ...
)

Также задается стандартная конфигурацая для модели на 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 ]
    }
]

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

Модели заданы. 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. А пока подписывайтесь на мой телеграм канал. Там рассказываю про нейронки с упором в сервинг.