Сервинг модели Grounding DINO с BentoML

Готовь сани летом, а план доставки ML-модели конечным пользователям — еще на этапе разработки. Иначе даже самая крутая обученная система будет пылиться без дела, а большинство пользователей о ней даже не узнает.

Запуск ML-моделей в production-среде — это тот еще квест:  

• необходимо хранить и версионировать модели;

• нужен мониторинг ресурсопотребления, производительности и дрейфа данных в реальном времени;

• нужна поддержка разнородных сценариев — пакетная обработка или стриминг.

К счастью, есть специальные инструменты, которые существенно облегчают такую задачу. В этой статье рассмотрим сервинг модели Grounding DINO при помощи популярного фреймворка BentoML. Под катом вас ждет много кода, скриншоты и практические советы.

P. S. Первая версия статьи написана для блога DeepSchool ^[1]. На Хабре публикуется с изменениями и дополнениями.

---

Кратко о BentoML

Для начала рассмотрим в самых общих чертах, что такое BentoML и с чем его едят. Это open-source-инструмент, который помогает:

1. оборачивать модель в сервис и обращаться к нему по API;

2. масштабировать сервис, а также объединять несколько сервисов в пайплайны;

3. мониторить сервис, используя Prometheus и Grafana;

4. хранить и версионировать модели.

Инструмент написан на Python — в этом его сила и слабость одновременно. Например, BentoML уступает Triton по скорости работы, но первым проще пользоваться. Другой козырь BentoML — интеграции с разными библиотеками: Hugging Face, scikit-learn, PyTorch, Gradio, MLflow и многими другими.

В платном инструменте BentoCloud есть возможность деплоить сервис одной командой и отслеживать показатели модели в UI. Также у него существует бесплатная версия для деплоя в кубере Yatai ^[2]. Правда, эта версия пока сыровата, и пользоваться ей сложнее.  

Тем не менее у BentoML множество примеров деплоя ^[3], в том числе и для LLM. Также стоит отметить подробную документацию ^[4] и обширные сообщества в Slack ^[5] и блог ^[6]. Всё это помогает пользователям лучше разобраться в нюансах применения инструмента.

Пару слов о Grounding DINO

«Минутка рекламы» BentoML закончена. Теперь о том, что за зверь этот Grounding DINO. Модель подробно описана в статье “Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection” ^[7]. По сути, это «сборная солянка» из нескольких довольно сложных моделей. Главная ее фишка — технология Open-Set Object Detection (OSOD), то есть детектирование объектов, чьих классов нет в обучающей выборке.

Скажем, нам нужно детектировать слонов на мировой черепахе как один объект. «Классические» closed-set детекторы вроде Faster R-CNN или YOLO так не умеют: они детектируют ровно то, что было размечено при обучении ^[8]. Можно заранее разметить «слон на черепахе» как отдельный класс и обучить модель сразу на него — тогда всё заработает. Но изменить логику ^[9] детектирования после обучения уже не выйдет: как разметили, так и будет находить. 

Grounding DINO ^[10] вместо таких танцев с бубнами добавляет еще одну модальность — язык (это и есть тот самый Grounding). Модель обрабатывает изображение и текстовый запрос совместно, обучаясь выравнивать визуальные и текстовые представления в общем семантическом пространстве. Благодаря этому детектор привязывается не к жёсткому списку классов из датасета, а к смыслу текстового запроса — и находит то, что вы описали словами, даже если такого класса при обучении не было. Поэтому «слон на черепахе как один объект» для нее уже не проблема.

С инструментами познакомились — теперь можно переходить к самому главному.

Собственно к сервингу

Не станем пускаться с места в карьер и разберем в статье простейший вариант такого сервинга, включая запуск и использование созданного сервиса. Весь код статьи и инструкция по его запуску доступны в репозитории ^[11]. Пример инференса модели Grounding DINO в dev-окружении на одной картинке выглядит следующим образом:

import requests

import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
from utils.utils import draw_detections # метод для отрисовки ббоксов на изображении

model_id = "IDEA-Research/grounding-dino-tiny"
device = "cuda:1" if torch.cuda.is_available() else "cpu" # если 1 гпу, замените на cuda:0

processor = AutoProcessor.from_pretrained(model_id)
model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id).to(device)

image_url = "<http://images.cocodataset.org/val2017/000000039769.jpg>"
image = Image.open(requests.get(image_url, stream=True).raw)
# Check for cats and remote controls
text_labels = [["a cat", "a remote control"]]

inputs = processor(images=image, text=text_labels, return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)

results = processor.post_process_grounded_object_detection(
    outputs,
    threshold=0.4,
    text_threshold=0.3,
    target_sizes=[(image.height, image.width)]
)
# Retrieve the first image result
print(results)
result = results[0]
for box, score, text_label in zip(result["boxes"], result["scores"], result["text_labels"]):
    box = [round(x, 2) for x in box.tolist()]
    print(f"Detected {text_label} with confidence {round(score.item(), 3)} at location {box}")

# Результат запуска кода
# Detected a cat with confidence 0.479 at location [344.7, 23.11, 637.18, 374.27]
# Detected a cat with confidence 0.438 at location [12.27, 51.91, 316.86, 472.43]
# Detected a remote control with confidence 0.476 at location [38.59, 70.01, 176.78, 118.17]

Для запуска кода выше необходимо создать виртуальное окружение и установить зависимости согласно инструкции в README.md.

Чтобы обернуть этот код в сервис с помощью BentoML, нужно выполнить следующие шаги:

1. определить конфигурацию для сборки сервиса;

2. создать класс GroundingDinoService. В нем будет вся логика сервиса: инференс модели и отрисовка результатов на изображении;

3. добавить валидацию входных параметров.

В итоге мы получим Bento — артефакт, предназначенный для развертывания ML-сервиса и включающий исходный код сервиса, все необходимые зависимости и сохраненные модели. С его помощью создается Docker-образ с автоматически сгенерированным API-сервером.

Запущенный сервис будет работать через HTTP API. Он начнет принимать изображение, промпт и пороги для инференса. В ответ сервис станет выдавать класс объекта на изображении, координаты bounding boxes и конфиденс.

Вся логика сервиса реализуется в файле service.py.

Определение конфигурации сборки сервиса

BentoML создает докерезированные сервисы. Dockerfile генерируется автоматически. Можно воспользоваться этой дефолтной версией или создать свой кастомный Dockerfile. Чтобы получить докер-контейнер, нужно создать Bento-объект.

Определим кастомную конфигурацию сборки нашего сервиса:

# образ по умолчанию при использовании gpu debian: "image": "nvidia/cuda:{spec_version}-cudnn8-runtime-ubuntu20.04"
runtime_image = bentoml.images.Image(
    python_version="3.11"
).pyproject_toml("pyproject.toml")

В этой части кода указывается базовый Docker-образ, версия Питона, дистрибутив, зависимости и всё, что обычно указывается в Dockerfile.

Больше интересной информации про настройку runtime_image можно найти в документации (раздел Define the runtime environment ^[12]).

Определение класса сервиса

При определении сервисов BentoML применяется классовый подход. Каждый класс — это пролетариат или буржуазия отдельный сервис, который может выполнять определенные задачи: предварительную обработку данных или инференс. Класс, указывающий на сервис BentoML, оборачивается декоратором @bentoml.service. Сервис может предоставлять доступ к одному или нескольким API-интерфейсам через HTTP.

Основные параметры сервиса определяются при помощи декоратора @bentoml.service для соответствующего класса:

@bentoml.service(
   name='grounding-dino-service', # имя сервиса
   traffic={'timeout': 300}, # допустимое время ответа сервиса
   resources={'gpu': 1}, # количество GPU или CPU, если GPU нет
   workers=1, # количество воркеров (экземпляров сервиса)
   image=runtime_image # кастомная конфигурация сборки сервиса
)
class GroundingDinoService:

    hf_model = bentoml.models.HuggingFaceModel(MODEL_ID)

    def __init__(self) -> None:
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModelForZeroShotObjectDetection.from_pretrained(self.hf_model).to(self.device)
        self.processor = AutoProcessor.from_pretrained(self.hf_model)
        print("Model grounding-dino loaded", "device:", self.device)

Поля, доступные для конфигурации сервиса, находятся в разделе документации Сonfigurations ^[13]. Чтобы не утонуть в множащихся параметрах, можно вынести определение сервиса в yaml-файл bentoml_configuration.yaml и перезаписать нужное из декоратора.

Модель в сервисе задается как переменная класса — hf_model, которая скачивается средствами BentoML. Таким образом, однократно загружаются веса модели при сборке Docker-образа. Параллельно выполняется оптимизированное скачивание весов. Если скачивать веса модели в конструкторе, это приведет к ошибке ^[14] model NotFound. Bento попросту не сможет создать ссылку на модель, привязанную к конкретному сервису.

В конструктор класса переносится уже инициализация процессора и модели.

Реализация API

С определением классов в общих чертах разобрались. Теперь самое время реализовать ключевые методы в классе GroundingDinoService:

1. _detect — приватный метод, содержащий основную логику инференса модели. Здесь находится вся логика из начального кода детекции объектов;

2. detect_image — публичный метод, представляющий собой обертку над _detect, которая возвращает результаты детекции;

3. render — публичный метод, который возвращает изображение с отрисованными результатами метода _detect.

Все методы обрабатывают одни и те же параметры — изображение и характеристики для инференса.

  @bentoml.api
  def detect_image(
      self,
      image: PILImage.Image,
      params: DetectionParams
      ) -> tp.List[tp.Dict[str, tp.Any]]:
      '''
      Detect objects in the image.
      '''
      return self._detect(image, params)

  def _detect(
      self,
      image: PILImage.Image,
      params: DetectionParams
  ) -> tp.List[tp.Dict[str, tp.Any]]:
      text_labels = params.detection_prompt
      inputs = self.processor(images=[image], text=text_labels, return_tensors="pt").to(self.device)
      with torch.no_grad():
          outputs = self.model(**inputs)
      results = self.processor.post_process_grounded_object_detection(
          outputs,
          threshold=params.box_threshold,
          text_threshold=params.text_threshold,
          target_sizes=[(image.height, image.width)],
      )
      return serialize_detections(results) 
 
  @bentoml.api
  def render(
      self,
      image: PILImage.Image,
      params: DetectionParams,
  ) -> PILImage.Image:
      '''
      Render detections on the image.
      '''
      result = self._detect(image, params)[0]
      image = draw_detections(image, result) 
      Path("images").mkdir(exist_ok=True)
      image.save("images/out_render.jpg")
      return image

Реализация вспомогательных функций serialize_detections ^[15] и draw_detections ^[16] доступна в исходниках репозитория на GitHub.

Методы detect_image и render помечаются декоратором @bentoml.api. Так легким движением руки эти методы автоматически превращаются в HTTP endpoints нашего сервиса. Теперь их можно вызывать извне через REST API. Подробнее про создание API с BentoML в разделе Service APIs ^[17] документации.

Валидация входных параметров сервиса

Как же отделить зерна от плевел и проверить параметры, которые подаются на вход в BentoML? Для этого поддерживается использование Pydantic ^[18]. Он позволяет описать классы-модели для валидации входных параметров сервиса.

Сама модель для такой валидации выглядит следующим образом:

class DetectionParams(BaseModel):
    detection_prompt: tp.List[tp.List[str]] = Field(..., description="List of lists of labels, e.g. [['a cat', 'a remote control']]")
    box_threshold: float = Field(default=0.25, description="Box threshold between 0 and 1")
    text_threshold: float = Field(default=0.25, description="Text threshold between 0 and 1")
    
    class Config:
        arbitrary_types_allowed = True.

Запуск сервиса

И вот наступает момент истины — запуск сервисов на BentoML. Здесь мы встаем перед развилкой: запускаться локально либо в Docker. Локальный запуск удобен, когда нужно протестировать сервис, внести изменение и не ждать каждый раз сборки. Применение Docker подходит для работы в production-среде.

Запуск сервиса локально для тестирования приложения

Для начала необходимо взять шаманский бубен выполнить в терминале следующую команду:

 bentoml serve --port 3025 --reload True

Бинго! Команда выполнена, и по ссылке http://localhost:3025/ ^[19] в браузере открывается Swagger UI.

По умолчанию сервис запускается на порту 3000. Однако порт можно поменять, скажем, на 3025 или любой другой с помощью флага --port , как это сделали мы в примере выше. Не забудем «поднять» еще один флаг --reload в значении True. Теперь при корректировке кода сервис автоматически принимает изменения без перезапуска.

Остальные параметры можно посмотреть в документации (раздел bentoml serve ^[20]).

Запуск сервиса c Docker

Теперь пойдем по второму пути и запустим сервис в Docker. Для этого сначала соберем Bento, выполнив в терминале команду bentoml build.

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

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

Сборка же образа сервиса выполняется следующей командой в терминале:

bentoml containerize grounding-dino-service:latest

В свою очередь, запуска сервиса на GPU выполняется такой командой:

docker run --rm --gpus '"device=1"' -p 3025:3000 grounding-dino-service:k76plrt6x6jkulue

Готово! Образы и контейнеры можно посмотреть стандартным способом через Docker.

Запущенный сервис доступен по адресу http://localhost:3025.

Использование сервиса

Запустить сервис — это только полдела. Отправлять запросы к нему можно при помощи кода, любого http client или Swagger UI, который генерируется в BentoML при запуске сервиса.

Согласно Swagger, в сервисе доступно два endpoint:

• /detect_image — возвращает json с классами, ббоксами и конфиденсами;

• /render — отрисовывает боксы на картинке, возвращает и сохраняет ее в папке images.

Рассмотрим каждый из этих двух вариантов подробнее.

Пример использования endpoint /render

Нажимаем кнопку Try it out, загружаем фото и вводим параметры:

Изображение с котиками в репозитории ^[21]

Нажимаем Execute и видим такой результат:

Аналогичный запрос можно выполнить через консольную утилиту cURL, получив изображение с отрисованными ббоксами в папке images:

curl -X 'POST' \
  '<http://localhost:3025/render>' \
  -H 'accept: image/*' \
  -H 'Content-Type: multipart/form-data' \
  -F 'image=@images/original_cat.jpg;type=image/jpeg' \
  -F 'params={
  "detection_prompt": [
    [
      "a cat", "a remote control"
    ]
  ],
  "box_threshold": 0.25,
  "text_threshold": 0.25
};type=application/json' \
  --output images/result_render.jpg

Пример использования endpoint /detect_image

Теперь перейдем ко второму детекту. Отправим запрос на endpoint /detect_image через SDK. Код доступен в файле client.py репозитория:

import bentoml
import requests
from io import BytesIO
from PIL import Image as PILImage

img_url = "<http://images.cocodataset.org/val2017/000000039769.jpg>"
img = PILImage.open(BytesIO(requests.get(img_url, stream=True).content))

with bentoml.SyncHTTPClient("<http://localhost:3025>") as client:
    result = client.detect_image(
        image=img,
        params={
            "detection_prompt": [["a cat", "a remote control"]],
            "box_threshold": 0.25,
            "text_threshold": 0.25,
        },
    )
    print(result)
    
    
# Результат в консоли
#[{'boxes': [[344.69305419921875, 23.10898208618164, 637.1846923828125, 374.27471923828125], [12.265024185180664, 51.91496276855469, 316.8591003417969, 472.438720703125], [38.58332061767578, 70.0059585571289, 176.77804565429688, 118.17623901367188], [332.17315673828125, 74.5613784790039, 370.69976806640625, 186.94830322265625]],
# 'scores': [0.4784787893295288, 0.4381333887577057, 0.475873202085495, 0.3312825560569763],
# 'text_labels': ['a cat', 'a cat', 'a remote control', 'a remote control']}]

Итого

В этой статье мы не стали вдаваться в сложные «игры разума» с кодом, а рассмотрели процесс создания простого сервиса на BentoML. Если обобщить, для этого нужно выполнить несколько шагов:

1. обернуть код инференса в класс;

2. добавить валидацию входных параметров через Pydantic;

3. обернуть методы класса декоратором @bentoml.api;

4. обернуть реализованный класс декоратором @bentoml.service c нужными параметрами.

Магия в том, что всё это делается в одном Python-файле менее чем на 100 строк!

Простота — это отлично, однако у BentoML есть и другие полезные функции для более сложных сценариев:

• версионирование и хранение моделей с помощью Bento локально и на S3;

• сборка и деплой сервиса с gitlab CI/CD;

• синхронность и асинхронность запросов;

• асинхронный инференс;

• онлайн- и офлайн-батчинг и многое другое.

Если вам интересно узнать об обработке таких сценариев при помощи BentoML в следующих статьях, не стесняйтесь и пишите об этом в комментариях.