В 9 раз быстрее и в 13 раз компактнее по сравнению со стандартными реализациями

Обучение ML‑модели — это только первый шаг к решению бизнес‑задачи. Далее необходимо создать эффективный механизм для развертывания модели в производственной среде и разработать serving‑стратегию, которая сможет масштабироваться в соответствии с текущим спросом.

В этой статье мы рассмотрим различные model serving стратегии и узнаем о технологиях, способных значительно повысить их эффективность. Мы рассмотрим три варианта организации model serving системы и сравним их производительность. Наша реализация будет ориентирована на инференс с помощью ЦП, но те же самые концепции могут быть применены и к ГП, поскольку предлагаемые здесь технологии (ONNX Runtime) поддерживают различные аппаратные платформы, включая графические и нейропроцессоры.

Весь исходный код можно найти в моем репозитории model‑serving. Тем, кого не интересует техническая сторона, я рекомендую сразу перейти к разделам «Результаты тестирования» и «Выводы».

Техническая справка

Прежде чем мы перейдем к примерам реализации, давайте сначала рассмотрим пару технических концепций: Open Neural Network Exchange (ONNX) и ONNX Runtime.

Open Neural Network Exchange

ONNX — это спецификация¹ (стандартный формат), предназначенная для представления моделей машинного обучения в виде графов вычислений (computational graphs), тем самым являясь общим языком для различных аппаратных платформ. Она определяет необходимые операции (операторы), типы данных и методы сериализации (с использованием Protocol Buffers), чтобы обеспечить совместимость и простоту развертывания моделей в различных средах. Спецификация ONNX поддерживает расширяемость за счет пользовательских операторов и функций и включает инструменты для визуализации модели, хранения метаданных и т. д.

ONNX Runtime

ONNX Runtime — это высокопроизводительный движок инференса (inference engine), предназначенный для эффективного выполнения моделей машинного обучения в формате ONNX на различных аппаратных платформах. Он служит кроссплатформенным ускорителем, который позволяет разработчикам внедрять модели, обученные в различных фреймворках, таких как PyTorch, TensorFlow, Keras и scikit‑learn, в производственные среды с минимальными накладными расходами.

Одним из главных преимуществ ONNX Runtime является его гибкая архитектура, которая поддерживает kernel‑based и runtime‑based Execution Provider»ы. Провайдеры на основе ядра (kernel‑based) реализуют определенные ONNX‑операции, оптимизированные для определенного оборудования (например, ЦП с CPUExecutionProvider, ГП с CUDAExecutionProvider), в то время как рантайм‑провайдеры могут выполнять целые графы вычислений или их части, используя специализированные ускорители, такие как TensorRT или nGraph.

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

Контекст задачи

Разобравшись с технической стороной вопроса, давайте перейдем к практическому применению: serving»у модели машинного обучения в производственной среде. Мы будем в некоторой степени опираться на предыдущую статью «ML Training Pipelines», где мы разработали модель для прогнозирования погодных условий. Вкратце, мы доработали модель MobileNet V3-small, которая включает в себя около 1,53 миллиона параметров и способна распознавать 11 различных паттернов синоптических карт.

После завершения обучения модели, следующий шаг — обеспечить ее эффективный serving в производственной среде. Чтобы упростить интеграцию с serving‑приложением, мы можем внести несколько улучшений в сам конвейер обучения.

Добавление преобразований входных данных в граф модели

В предыдущей статье мы обсуждали, как сохранить модели PyTorch и ONNX в виде артефактов в конвейерах Kubeflow для последующего использования или непосредственного развертывания в производственной среде. Однако существует один полезный способ улучшить этот подход — встраивание преобразований изображений непосредственно в граф вычислений модели. Это даст нам два больших преимущества:

Модульность и упрощение: Включение преобразований входных данных в граф модели позволяет отделить логику обработки входных данных от логики serving»а. Это делает конфигурацию более модульной и упрощает интеграцию. Кроме того, это сводит к минимуму сторонние зависимости в части serving»а, что приводит к более компактным образам Docker и более быстрому запуску.
Оптимизированная скорость обработки: Если преобразования входных данных будут встроены в граф, то ONNX Runtime сможет оптимизировать их, что еще больше повышает общую скорость обработки запросов.

Чтобы реализовать это улучшение, нам необходимо понять, какие преобразования используются в MobileNet_V3_Small_Weights.DEFAULT.transformes():

ImageClassification(
    crop_size=[224]
    resize_size=[256]
    mean=[0.485, 0.456, 0.406]
    std=[0.229, 0.224, 0.225]
    interpolation=InterpolationMode.BILINEAR
)

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

..

class ModelWithTransforms(Module):  # type: ignore[misc]
    def init(self, model: MobileNetV3) -> None:
        super(ModelWithTransforms, self).__init__()
        self.model = model
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
        self.register_buffer("targ_h", torch.tensor(224))
        self.register_buffer("targ_w", torch.tensor(224))

    def transform(self, img: torch.Tensor) -> torch.Tensor:
        # При необходимости добавьте размерность батча.
        if img.dim() == 3:
            img = img.unsqueeze(0)
        resized = F.interpolate(img, size=256, mode="bilinear", align_corners=False)
        , , curr_h, curr_w = resized.shape
        pad_h = torch.clamp(self.targ_h - curr_h, min=0)
        pad_w = torch.clamp(self.targ_w - curr_w, min=0)
        padding = [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]
        
        padded = pad(resized, padding)
        start_h = torch.clamp((curr_h + pad_h - self.targ_h) // 2, min=0)
        start_w = torch.clamp((curr_w + pad_w - self.targ_w) // 2, min=0)
        cropped = padded[..., start_h : start_h + self.targ_h, start_w : start_w + self.targ_w]
        normalized = (cropped - self.mean.to(cropped.device)) / self.std.to(cropped.device)
        return normalized

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.transform(x)
        return self.model(x)

model обозначает исходную обученную модель. Чтобы сохранить новую модель с интегрированными преобразованиями:

..

   model_with_transform = ModelWithTransforms(model)
    model_with_transform.to(device)
    torch.onnx.export(
        model_with_transform,
        model_input,
        f"{onnx_with_transform_model.path}.onnx",
        opset_version=opset_version,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size", 2: "height", 3: "width"},  # Динамический размер батча, высота и ширина
            "output": {0: "batch_size"},  # Динамический размер батча для вывода
        },
    )

Оффлайн-оптимизация ONNX-графов

Как гласит официальная документация ONNX Runtime:

Все оптимизации можно выполнять как в онлайн, так и в офлайн‑режиме. В онлайн‑режиме при инициализации сеанса инференса мы также применяем все доступные оптимизации графа перед непосредственным выполнением инференса модели. Однако применение всех оптимизаций при каждом запуске сеанса может увеличить время запуска модели (особенно для сложных моделей), что может быть критично в производственных сценариях. Именно здесь на помощь приходит оффлайн‑режим. В этом режиме ONNX Runtime после оптимизации графа сериализует полученную модель на диск. Впоследствии мы можем сократить время запуска, отключив все онлайн‑оптимизации и используя уже оптимизированную модель.

В зависимости от размера модели, такая оптимизация может существенно сократить время запуска инстанса, что, в свою очередь, позволит быстрее масштабировать его в производственных системах при высоких нагрузках. Реализация очень проста. Все, что необходимо сделать, — это добавить небольшой компонент в исходный конвейер, который будет принимать ONNX‑модель с преобразованиями в качестве входных данных. Вот пример того, как может выглядеть реализация:

from kfp.dsl import Input, Metrics, Model, Output


def onnx_optimize(
    onnx_with_transform_model: Input[Model],
    optimization_metrics: Output[Metrics],
    optimized_onnx_with_transform_model: Output[Model]
) -> None:
    import time
    import onnxruntime as rt

    start_time = time.time()
    sess_options = rt.SessionOptions()
    sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
    sess_options.optimized_model_filepath = optimized_onnx_with_transform_model.path
    rt.InferenceSession(f"{onnx_with_transform_model.path}.onnx", sess_options)
    optimized_onnx_with_transform_model.framework = (
        f"onnxruntime-{rt.__version__}, graphOptimizationLevel-{str(sess_options.graph_optimization_level)}"
    )
    optimization_metrics.log_metric("timeTakenSeconds", round(time.time() - start_time, 2))

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

При работе в оффлайн‑режиме убедитесь, что используются те же параметры (например, execution provider»ы, уровень оптимизации) и аппаратное обеспечение, что и на целевой машине, на которой будет выполняться инференс модели. Например, вы не сможете запустить модель, предварительно оптимизированную для выполнения на ГП, на машине, оснащенной только ЦП.
Если включена оптимизация топологии, оффлайн‑режим может использоваться только на оборудовании, совместимом со средой, в которой производилась оффлайн‑оптимизация. Например, если топология модели оптимизирована под AVX2, для оффлайн‑модели потребуются процессоры, поддерживающие AVX2.

Model Serving стратегии

После того как оптимизированная модель готова, мы можем приступить к работе над serving‑приложением. В этой статье мы проведем сравнительный анализ производительности трех различных serving‑стратегий:

Наивный model serving на PyTorch и FastAPI (Python): В этой конфигурации используется PyTorch с включенными model.eval() и torch.inference_mode(). Никакие оптимизации ONNX или ONNX Runtime не применяются. Вместо этого мы реализуем serving модели непосредственно из ее сохраненного state_dict после обучения. Хотя этот подход наименее оптимален, на практике он довольно распространен (с возможными альтернативами FastAPI в виде Flask или Django), что делает его хорошим базисом для наших тестов.
Оптимизированный model serving на ONNX Runtime и FastAPI (Python): В рамках этого подхода для serving»а модели мы используем ONNX Runtime. Логика преобразования входных данных встроена непосредственно в граф вычислений модели, для которого выполняет оффлайн‑оптимизация, что обеспечивает более эффективную альтернативу по сравнению с наивным подходом.
Оптимизированный model serving на ONNX Runtime и Actix‑Web (Rust): Здесь для serving»а мы используем конфигурацию на основе Rust с ONNX Runtime (собранного из исходного кода и использующего оболочку pykeio/ort) и Actix‑Web. Как и в предыдущей конфигурации, логика преобразования входных данных встроена в граф модели, и для достижения максимальной производительности применяется оффайн‑оптимизация.

Бенчмарк

При интерпретации результатов бенчмарков следует понимать, что они не могут рассматриваться как универсальные значения, поскольку абсолютная производительность может значительно варьироваться в зависимости от оборудования, операционных систем и реализаций стандартных библиотек C (например, glibc или musl), которые влияют на двоичный интерфейс приложения (ABI).

Кроме того, производительность может варьироваться в зависимости от размера входных изображений. Поэтому для производственной среды важно иметь представление о размерах изображений. В данном примере основное внимание уделяется относительным различиям в производительности между различными serving‑стратегиями.

Самый надежный способ оценить производительность службы модели на конкретной хост‑системе — это провести прямое тестирование в этой среде.

Хост-система

Аппаратное обеспечение: Apple M2 Max
Операционная система: macOS 15.0.1
Docker:
- Engine v27.2.0
- Desktop 4.34.3

Контейнеры

Распределение ЦП: Каждому контейнеру было выделено по 4 ядра ЦП.
Распределение памяти: Память распределялась динамически, предоставляя каждому контейнеру столько памяти, сколько ему требовалось.
Конфигурация воркеров и потоко: Для обеспечения максимального использования выделенного ЦП в каждом контейнере — до 400%, что соответствует 4 ядрам процессора — был тщательно отслежен и предотвращен перерасход ЦП. Чтобы достичь наилучшей производительности, были реализованы следующие конфигурации:
- onnx_serving:
  - Uvicorn Workers: 4
  - ONNX Runtime Session Threads:
    - Intra‑Op Threads: 1
    - Inter‑Op Threads: 1
- torch_serving:
  - Uvicorn Workers: 4
- rust_onnx_serving:
  - Actix Web Workers: 4
  - ONNX Runtime Session Threads:
    - Intra‑Op Threads: 3
    - Inter‑Op Threads: 1

Конфигурация бенчмарка

Инструмент бенчмаркинга: apache benchmark.

ab ‑n 40 000 ‑c 50 ‑p images/rime_5868.json ‑T 'application/json' ‑s 3600 “http://localhost:$port/predict/”

‑n 40 000: всего 40 000 запросов.
‑c 50: параметр, определяющий степень параллелизма (concurrency)
Изображение в пейлоаде: images/rime_5868.jpg:

Размер оригинала: 393 КБ.
Размер пейлоада после PIL‑сжатия и кодирования в base64 (увеличение на ~33%): 304 КБ.

Реализации

Поскольку model serving предполагает большое количество кода, я решил предоставить ссылки на соответствующие каталоги в репозитории GitHub. Такой подход делает статью более читабельной и компактной, позволяя читателям легко просматривать код с подсветкой синтаксиса на GitHub и лучше понять структуру проекта.

Наивный Model Serving на PyTorch/FastAPI

Model Serving на ONNX-Runtime/FastAPI (Python)

Model Serving на ONNX-Runtime/Actix Web (Rust)

Результаты бенчмарка

Показатели производительности

Показатели развертывания

Выводы

ONNX Runtime значительно повышает производительность: Преобразование моделей в ONNX и их serving с помощью ONNX Runtime значительно повышает пропускную способность и сокращает задержки по сравнению с serving»ом с помощью PyTorch. В частности:
- onnx‑serving (Python) обрабатывает примерно в 7.18 раза больше запросов в секунду, чем torch‑serving (255.53 против 35.62 запросов в секунду).
- rust‑onnx‑serving (Rust) обеспечивает примерно в 9.23 раз более высокую пропускную способность, чем torch‑serving (328.94 против 35.62 запросов в секунду).
Реализация на Rust обеспечивает наилучшую производительность: Несмотря на более высокое потребление памяти, чем при использовании Python ONNX serving»а, Rust‑реализация демонстрирует более высокую производительность и предоставляет дополнительные преимущества в размере развертывания и времени запуска. Вот некоторые из них:
- Пропускная способность: rust‑onnx‑serving примерно в 1.29 раза быстрее, чем onnx‑serving (328.94 против 255.53 запросов в секунду).
- Время запуска: Приложение Rust запускается за 0.348 секунды, что более чем в 12 раз быстрее, чем torch‑serving (4.342 секунды) и примерно в 4 раза быстрее, чем onnx‑serving (1.396 секунды).
- Размер образа Docker: Размер образа Rust‑конфигурации составляет 48.3 МБ, что примерно в 13 раз меньше, чем у torch‑serving (650 МБ) и примерно в 6 раз меньше, чем у onnx‑serving (296 МБ).
Разница в потреблении памяти: Более высокое потребление памяти в Rust по сравнению с Python ONNX связано с различиями в реализациях и используемых библиотеках:
- Разница в обработке изображений: Реализация на Rust использует менее оптимизированные алгоритмы обработки изображений по сравнению с популярными библиотеками Python, такими как PIL и NumPy, что приводит к увеличению объема потребляемой памяти.
- Эффективность библиотеки: Крейт ort Rust является неофициальной оболочкой и может управлять памятью иначе по сравнению с официальным ONNX Runtime SDK для Python, который является достаточно зрелым и высоко оптимизированным.
- Конфигурация потоков: Реализация на Rust задействует больше intra‑opt потоков, что приводит к незначительному увеличению потребления памяти. Однако это не является основной причиной наблюдаемой разницы в потреблении ресурсов.

Но важнее другое: экосистема машинного обучения Python зрелая и обширная. Если переписать serving‑стратегии на Rust, это может повлечь за собой определенные трудности. Например, потребуются дополнительные усилия по разработке, могут возникнуть потенциальные компромиссы в производительности, если оптимизированные библиотеки не будут доступны или их придется создавать самостоятельно, а также добавятся другие сложности. Тем не менее, преимущества Rust могут быть настолько значительными, что оправдают все эти дополнительные усилия, особенно если это соответствует бизнес‑требованиям.

Использование решений, оптимизированных для инференса, таких как ONNX Runtime, может значительно повысить производительность model serving»а, особенно для более крупных моделей. Хотя в данной статье рассматривается небольшая модель (MobileNet V3-small, с примерно 1,53 миллионами параметров), преимущества ONNX Runtime становятся более заметными при работе с более сложными архитектурами. ONNX Runtime способен оптимизировать графы вычислений и эффективно использовать ресурсы, что приводит к увеличению пропускной способности и сокращению задержек. Это делает его незаменимым выбором для масштабирования modelserving решений.

1
В различных источниках ONNX может упоминаться как промежуточное представление моделей.
2
При реализации подобных преобразований всегда важно убедиться, что они выполняются идентично исходному преобразованию, которое вы копируете. Полезным шагом проверки может стать сравнение изображений до и после обоих преобразований.

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

Автоматизация ML‑экспериментов с помощью GitLab CI/CD и CML — 17 июня в 20:00

Если же вы хотите системно прокачаться в профессии, рекомендуем обратить внимание на курс «MLOps», который стартует 26 июня.

Model Serving в 9 раз быстрее! И никакой замены оборудования