Как мы от torch моделей уходили (и ещё не ушли)

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

Сейчас EdTech-компаниям как никогда важно делать упор на постоянное совершенствование и оптимизацию технологических решений. Так, одно из ключевых направлений работы системы «Антиплагиат» — это обнаружение фрагментов текста, созданных при помощи ИИ. Мы понимаем, что в такой динамичной области мы должны развиваться вместе с новыми технологиями. Для этого нужно, чтобы наша информация была организована таким образом, чтобы мы могли быстро адаптироваться к изменениям и отвечать на потребности наших пользователей, а не копаться в бесконечных фрагментах разрозненного кода.

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

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

Мы, команда исследователей компании Антиплагиат, ещё очень давно начали продумывать шаги для кардинального изменения ситуации — занялись унификацией кодовой базы. Так были созданы библиотека утилит, где собирались основные утилиты, использующиеся во всех сервисах, шаблоны создания сервисов, проторепозиториев и многое-многое другое. Но из-за высокой загруженности людей, поставленных на данную задачу, процесс рефакторинга встал на паузу.

Когда за задачу взялись снова, в глаза сразу бросились рудименты разработки прямо на проде. Такими рудиментами, например, являются модели, написанные с использованием PyTorch и встроенные с помощью него же в рабочий пайплайн.

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

Так, примером с полей может послужить размер образа сервиса машинной генерации. Когда в нём использовалась модель на PyTorch, образ весил больше 5 Гб, а после преобразования — 1,06 Гб, что показывает, насколько существенной может быть разница в использовании моделей на PyTorch и модели унифицированной, о которой мы говорим чуть дальше.

А в чём же тогда явные минусы подхода с хранением кода моделей в репозитории? Частично мы уже обозначили следующие:

• увеличение образов сервисов за счёт лишних зависимостей, тянущихся за моделями (Torch, TorchVision, Albumentations и др.) ;
• теоретическое замедление работы моделей, т. к. у простых торчевых моделей нет оптимизаций;
• разнородность использования моделей в сервисах, впоследствии проблематичность поддержки сервисов;
• наличие артефактов обучения моделей в промышленном репозитории: не так часто разработчики заморачиваются написанием отдельного инференса для моделей.

Всё это приводит к трудностям в поддержании моделей, их замене.

• перевод моделей в общий формат;
• разделение репозиториев на исследовательские и репозитории сервисов (можно ещё и третий придумать) . Исследовательские репозитории в данном контексте будут включать в себя эксперименты исследователей (код обучения, обработки данных, самой модели) , а репозитории сервисов — уже непосредственно кодовую базу сервисов, не включая код самих моделей.

Здесь мы собрались, чтобы обсудить, как прошёл перевод моделей к единому формату. Начнём с того, что до этого уже существовало решение, которое просто не было распространено на сервисы. Таким образом, встала задача:

• доработать существующее решение;
• начать распространение этого решения на существующие сервисы.

Прежде чем перейти к самим доработкам, рассмотрим, какие популярные варианты инференса моделей существуют на данный момент.

Одними из самых популярных вариантов являются следующие:

• ONNX;
• TorchScript;
• OpenVINO.

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

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

В качестве преимуществ такого формата для инференса моделей можно выделить:

• кросс-платформенность (мы можем создать модель на PyTorch, конвертировать ее в ONNX, а затем использовать как на Linux, так и на других платформах) ;
• уменьшение количества зависимостей (как правило, PyTorch несёт за собой довольно внушительное количество зависимостей, которое обычно не требуется при инференсе, для ONNX же обычно требуется один легковесный ONNX Runtime) ;
• открытый код проекта (проект дорабатывается в угоду пользователям в том числе и компаниями-гигантами, быстрее отлавливаются и чинятся баги) ;
• оптимизация времени работы модели.

TorchScript — специальный инструмент от PyTorch для сериализации и оптимизации моделей. Он также представляет модель в виде вычислительного графа, но, в отличие от ONNX, работает только с моделями из одного фреймворка — PyTorch, остальные же не поддерживаются.

Из преимуществ этого инструмента можно выделить:

• удобство конвертации моделей из PyTorch;
• возможность бесшовно перейти к эксплуатации на С++.

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

Преимущества инструмента:

· производительность;

• уменьшение количества зависимостей в проекте;
• уменьшение размеров моделей;
• наличие большого количества инструментов для моделей (Model Zoo, Deep Learning Model Optimizer, Post-training Optimization Tool и др.) .

А теперь можно перейти к делу.

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

• невозможность масштабируемости;
• узкая направленность на один движок — ONNX.

Если про первый недостаток всё понятно, то второй может быть не очень очевидным. Мы же хотим унифицировать модели, зачем нам нужно несколько инференсных движков? Ответом на этот вопрос может быть то, что разработчикам следует стремиться изучать новые технологии. Со временем будут появляться всё более универсальные и оптимизированные инференсные движки, которые необходимо будет внедрить в уже существующее решение. Таким образом, мы стремимся избавиться от лишних заморочек и сделать данный процесс безболезненным.

Доработка существующего решения явила собой:

• добавление базового класса для обёрток инференсных движков;
• добавление метода для гибкого выбора модели с помощью конфига;
• обновление существующего класса для ONNX, попытка сделать его более гибким к использованию.

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

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

Таким образом, решение было перенесено в сервисы:

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

В результате распространения решения у некоторых сервисов было выявлено увеличение производительности, упрощение кодовой базы (за счёт вырезания модели из репозитория и унификации инициализации и вызова модели) , уменьшение докер-образа (при условии изъятия Torch и ему сопутствующих библиотек и фреймворков) .

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

На самом деле у нынешнего решения тоже есть проблемы недостатки, представим некоторые из них:

• не все модели представляется возможным перевести в ONNX;
• сложности в переводе некоторых моделей на ONNX.

Решением этого будет добавление в библиотеку новых движков, которые будут либо упрощать перевод, либо делать его возможным, поэтому это ещё далеко не конец…