Нейросеть, сегментирующая изображение человека в кадре: как ускорить её в четыре раза?
Привет, Хабр! Это Дмитрий Балиев из SberDevices. В этой статье, написанной по докладу с конференции Saint HighLoad++ речь пойдёт о том, как мы обрабатываем алгоритмами видео в Web-браузерах для сервиса конференции SaluteJazz.
Расскажу про контекст и ограничения, сам движок и особенности его реализации. Раскрою тайну, как мы работаем с графами вычислений, как инферим нейросети, и как затем всё это собираем, оптимизируем и тестируем. В конце — несколько полезных советов, как делать нейросети удобнее для встраивания.
Тизер: мы взяли нейросеть, отвечающую за сегментацию человека. Она используется в таком эффекте, как замена фона. Запустили её в двух вариантах:
На актуальной версии ONNX Runtime в Web, причём используя один из самых передовых бэкендов WebGPU, который ещё мало где поддерживается.
На движке, который сами написали.
Оказалось, наша сеть может работать в четыре раза быстрее. А дальше о том, как удалось достичь таких результатов.
Продукт: для чего мы всё это делаем
Сервис видеоконференции SaluteJazz — это умная платформа для коммуникаций с использованием AI с 1 млн активных пользователей ежемесячно. Еженедельно на платформе проходит 300 тысяч видеоконференций.
Внутри у нас работают AI-алгоритмы:
Шумоподавление.
Обработка видео:
Замена фона;
Улучшение освещения;
Бьютификация.
Транскрибация речи.
Jazz XR-режим, позволяющий создавать виртуальные пространства для встреч.
Всё это мы запускаем на разных платформах:
Desktop App. Приложение написано на фреймворке Electron. По сути, это обёрнутый браузер с расширенными возможностями. Запускается на всех основных платформах — Windows, MacOS, Linux.
Mobile App для iOS и Android. В мобильных приложениях сейчас нет акцента на обработке видео, но нам нужно в любой момент иметь возможность вкатить туда нужные фичи.
Web. С видеоконференциями часто возникает ситуация, когда человеку прислали ссылку на встречу, а у него не установлено приложение. Но даже в этом случае мы всё равно хотим предоставить полноценный опыт — все фичи, которые доступны на платформе. Через браузер пользователь тоже может пользоваться заменой фона, улучшением картинки и так далее. В общем, web для нас — важная базовая платформа, в которой пока меньше всего возможностей. Поэтому в разработке мы отталкивались идеи: если сможем сделать для web, то и на десктопе, и на мобилках — тоже сможем.
Обработка видео
В любой платформе видеоконференции выполняется множество задач — обработка медиа, работа с кодеками, управление сетью и UI. Наша команда отвечает за процесс, который происходит после получения видео, стрима, кадров.
Мы делаем магию, обрабатываем кадры, и возвращаем результат. Он обрабатывается кодеками, отправляется на сервер, маршрутизируется и в итоге попадает к пользователям.
Рассмотрим пару примеров задач, которые мы решаем:
Бьютификация. Нужно найти на кадре лицо, определить на нём регион интереса, где кожа находится в bounding-боксе лица. Дальше наложить размытие и получить более приятную картинку.
Размытие фона. Это частный случай замены фона, самый сложный, в котором нам нужно сначала найти пиксель, относящийся к человеку, то есть решить задачу сегментации. Получив сегментационную маску, нужно обработать фон размытием и сделать композицию.
Естественно, у нас есть свои пожелания к этим процессам:
Работать в браузере. Как я уже сказал, для нас браузер — это базовая платформа, и мы хотим хорошо в нём работать.
Поддержка малой командой. Мы хотим писать максимально универсальный код и не поддерживать «зоопарк» решений для разных платформ.
Минимизация нагрузки. С одной стороны, это сделает наш продукт быстрее и приятнее в использовании. С другой, поскольку наши пользователи работают в основном на ноутбуках, то это вопрос продолжительности работы от батареи.
Работа в 30 FPS. Это стандартная скорость, с которой работают камеры. Но мы хотим не только этого, но и чтобы вся система работала быстро. Если мы возьмём весь вычислительный бюджет только на себя, то будет тормозить медийка, кодеки, UI — это неприятно.
Работа на слабом железе. Мы хотим хорошо работать даже на относительно слабом железе. Внутренний ориентир — процессоры Intel со встроенной графикой i3, i5 2018 года выпуска.
Правило №1: без лишних копирований данных
В самом начале разработки мы ввели правило: не делать лишних больших копирований данных, особенно между центральным процессором и графическим, потому что это медленно. Мы работаем с большими кадрами с разрешением 720p, 1080p, иногда даже 4k. Дополнительную проблему создаёт то, что графические API — обычно асинхронные. Если мы попытаемся передать кадр с центрального процессора на графический, столкнёмся с барьерами асинхронной работы, а это дополнительные «палки в колёса».
В идеальном мире мы получаем кадр с камеры, обрабатываем его на графическом процессоре и получаем результат, который уходит в кодеки. И нигде не переходим в память центрального процессора.
Компоненты
Исходя из этого, у нас есть чёткое представление, что именно мы хотим сделать, в каких ограничениях и какие компоненты для этого нужны:
Движок для построения пайплайнов — чтобы из разных кусочков, как из кубиков, собирать конечный эффект.
Движок для inference нейросети.
Интеграция и сборка.
Графовый движок
Начнём с движка для построения пайплайнов: почему мы на нём остановились, какие плюсы он даёт, и почему это важно.
Рассмотрим графовый движок на примере реализации одного из наших пайплайнов. В упрощенном виде посмотрим на пайплайн, который улучшает картинку.
Сначала мы уменьшаем размер изображения для нейросети. Затем запускаем нейросеть, и она находит на этой картинке лица. Так мы получаем местоположение лиц на изображении. Дальше по частям картинки считаем набор коэффициентов, после чего применяем конечную обработку кадра. Вроде бы, всё просто — такой пайплайн легко построить в коде. Но уже сложнее, если мы хотим запускать не один пайплайн, а сразу несколько параллельно — например, к улучшению картинки добавить бьютификацию.
В этом случае ко всем операциям для улучшения картинки, добавляются операции, необходимые для бьютификации. Для этого берём картинку, уменьшаем её размер для обработки нейросетью и снова ищем на картинке лица. На этих лицах находим регион интереса и применяем к нему сглаживание.
Очевидно, что в этих пайплайнах операции повторяются, а нам бы этого не хотелось. Часть из этих операций — тяжелая, в частности, запуск нейросетей. А ведь мы хотим минимизировать нагрузку на устройство пользователя. В идеале, чтобы Downscale и запуск нейросети поиска лиц происходил один раз, а дальше мы по одной ветке считали коэффициенты для улучшения картинки, а по другой — бьютификацию.
Такую систему несложно построить, но есть подводные камни. Например, если мы хотим использовать только один из этих эффектов — не улучшать картинку, но оставить бютификацию. Можно использовать грубый подход как на схеме ниже, но тогда будут лишние вычисления:
Чтобы это решить, пришли к архитектуре, которую я называю pool-архитектурой.
Предположим, мы хотим получить картинку, к которой применена бьютификация. Очевидно, результат будет зависеть от выхода последнего эффекта в ветке бьютификации. Пропишем эту зависимость. Нам нужно получить данные из последней ячейки. Чтобы её просчитать, нужны данные о том, как мы посчитали регионы интереса. После этого просчитываем следующую зависимость.
Так мы трассируем сетку. Чтобы получить результат (картинку, к которой применена бьютификация), нужно посчитать чёрные ноды. Серые мы запускать не будем, чтобы не грузить устройства пользователя лишними вычислениями.
Если же мы хотим использовать оба эффекта, то считаем результат зависимым от ветки улучшения картинки. Тогда нам придётся посчитать все ноды. А ещё сделать кэширование результатов нахождения лиц, чтобы не считать два раза тяжелую часть и минимизировать нагрузку.
Некоторые детали реализации, которые хочется упомянуть:
Все пайплайны инициализированы. Это позволяет быстро включать и выключать эффекты. Как только изменился граф зависимости, мы понимаем, что теперь на этом кадре нужно рассчитать дополнительную ноду и, естественно, у нас сразу всё заработает.
Считаем только необходимый подграф, тем самым минимизируем нагрузку.
Сами операции делаем как обычные функции и аккуратно оборачиваем.
Пайплайны собираются в коде.
Мы не стали делать построение пайплайнов с помощью конфигурационных файлов в этой системе. То, что пайплайны сами собираются в коде, удобно и даёт подстраховку в виде строгой типизации языка программирования, которая помогает допускать меньше ошибок.
Теперь мы можем собирать из наших кубиков пайплайны. Осталось решить, как мы будем делать кубики, которые относятся к inference нейросетей.
Inference нейросети
На большинстве платформ, для инференса нейросети уже есть инструменты, позволяющие делать это с лучшей производительностью. Например, у NVIDIA есть TensorRT, у Apple — CoreML, и так далее. Но мы не хотим поддерживать большой «зоопарк» решений. Да и не для всех платформ они есть. А ведь нам, кроме нейросетей, нужно писать ещё и обработку с помощью дополнительного фреймворка. Мы выбрали категорию решений, не заточенных на inference нейросетей, но которые, тем не менее, дают полноценный доступ к графической карте. Остановились мы на семействе API OpenGL.
Есть полноценный OpenGL для компьютеров и ноутбуков, OpenGL ES для мобильных платформ и WebGL для работы в браузере с графической картой. Если мы напишем код, совместимый с самым слабым API общего назначения из этой тройки, то он будет работать везде: на мобилках, десктопе, в Web, и можно делать на нём обработку изображения (даунскейлы, постпроцессинг).
Теперь расскажу о реализации inference на OpenGL и подобных языках.
Архитектурные идеи
При работе с графикой есть шейдерные программы, выполняющие математические операции. Они читают данные из входной и записывают результаты в выходную текстуру. Сами программы могут работать в N потоков. N зависит от того, какая конкретно графическая карта и количества ядер в ней. Выстраиваем цепочку таких программ, получаем нужный эффект, выполняем цепочку вычислений и получаем посчитанную нейросеть.
Важно отметить, что мы хотим минимизировать количество таких шейдерных программ. При работе с такой архитектурой именно операции с памятью — самые медленные. Мы не хотим большого числа операций чтения тензоров из текстуры и записи тензоров в текстуры. Чем их меньше, тем лучше.
Чтобы это сделать, нужно понять, какие операции мы хотим запускать в шейдерных программах.
Есть два класса операций для нейросетей:
«Толстые» — стандартные свёртки, poolings, upscales. Все эти операции характеризуются особенностью: одно выходное значение зависит сразу от нескольких значений на входе. В общем случае мы не можем собрать такие операции в одну шейдерную программу, их придётся разделять.
«Тонкие» — поэлементные операции. В нейросетях самый яркий представитель — это функция активации. В таких операциях одно значение на выходе соответствует одному значению на входе. Их можно комбинировать в бесконечное количество.
Базовая архитектурная идея такая:
1 шейдерная программа = 1 «толстая» операция + [0; N] «тонких» операций
В одну шейдерную программу берём одну «толстую» операцию, поверх неё ставим от 0 до N «тонких», пока не встретим следующую «толстую». Таким образом соединяем часть шейдерных программ вместе.
«Толстые» операции мы можем объединить, если пишем специализированный код. Для некоторых мест в архитектурах такой код стоит написать. Один из примеров — это места в структуре декодера модели для сегментации, в которых у нас происходит upscale, после которого сразу идёт свёртка.
Поскольку размер входной текстуры такого слоя в четыре раза меньше, чем выходной, получаем плюс от соединения этих операций. Для этого делаем операцию Upconv, которая совмещает сначала Upscale, а потом конволюцию.
Это даёт:
10% прирост производительности от совмещения «толстых» и «тонких» операций в одной программе. Причём это не просто 10% от inference одного из слоев или нейросети, а от всего пайплайна в целом — реальные 10% производительности в продукте.
15% прирост производительности от совмещения upscale и свёртки в одну операцию.
С таким решением мы сделали наш первый релиз.
Сборка
Теперь расскажу про сборку пайплайна. У нас было описание пайплайна модели (путь к модели, параметры конфига) в стандартном формате ONNX Runtime. Всё это отправляется в кодогенератор, который на выходе заполняет шаблоны на C++. В таких шаблонах у нас лежат заготовки кода на GLSL (языке шейдерных программ), а затем всё это попадает в стандартный компилятор языка C++.
Поскольку мы делаем это под Web, основной вариант сборки — WebAssembly, но помимо этого можем собирать и под десктопные платформы (MacOS, Linux, Windows). На выходе мы всё это встраиваем в продукт и неминуемо причиняем счастье пользователю.
Таким был первый релиз. В принципе, всё работало хорошо, но у некоторых пользователей эффекты подтормаживали. Встал вопрос, как оптимизировать решение.
Оптимизация
Расскажу про конкретную малоизвестную оптимизацию, давшую наибольший прирост в производительности.
Базовая проблема работы с WebGL как инструментом для inference нейросетей заключается в отсутствии в стандарте поддержки типа шейдерных программ, которые называются Compute Shaders. Они специально предназначены для математических вычислений, матричных операций и так далее. Дело в том, что в какой-то момент разработку стандарта WebGL заморозили и переключились на следующий стандарт WebGPU. Скоро его можно будет использовать для таких проектов, но пока его поддержка и степень зрелости оставляют желать лучшего.
В результате получилась ситуация, когда на десктопах и мобильных устройствах всё есть, а в Web нужных нам Compute Shaders — нет.
Как из этого выпутываться, нам подсказал блог-пост, который выпустил Google в 2022 году. Они тоже столкнулись с такой проблемой — согласно замерам, производительность WebGL без Computer Shaders достигала лишь 25% от полноценного OpenGL.
Но оказалось, часть потерянной производительности можно отыграть, если использовать подход Multiple Render Targets (далее MRT). Это подход, когда одной шейдерной программой мы пишем не в одну выходную текстуру, а в несколько параллельно.
Дело в том, что свёртка — это memory-bound-операция, а операции с памятью, например, чтение, занимают больше времени, чем вычисления в этой программе. Если сможем оптимизировать операции с памятью, свёртка будет работать быстрее.
В Google попробовали это применить, и выяснили, что с оптимизацией MRT производительность WebGL достигла 90% от изначальной производительности OpenGL. Проблема в том, что нет никаких деталей и примеров реализаций. Оставалось только попробовать самим.
Чтобы стало понятно, почему вообще эта идея работает, вспомним, как работает свёрточная операция.
Как работает свёртка
Рассмотрим супер простой пример. Есть маленький тензор, который мы сворачиваем с простым ядром 3×3 для получения результата.
Сначала берём окошко, делаем поэлементное перемножение, получаем результат. Так двигаемся по нашей матрице, смещаем окошко, получаем следующее число.
Четыре раза повторяем и получаем полный результат свёртки. Давайте это запомним.
Теперь будем считать операции с памятью, которые нужны для этого.
Без MRT. Мы не используем Multiple Render Targets оптимизацию и просто считаем одно значение на выходе. Для этого нужно прочитать из памяти 9 значений нашей свёртки 3×3 и 9 значений из нашего окошка — тоже 3×3. Итого, 18 операций чтения.
MRT в четыре пикселя. Если мы хотим посчитать на выходе одновременно четыре значения и записать их в четыре разные текстуры, то нужно прочитать данные для всех этих четырёх операций. Получаем девять значений ядра свёртки, поскольку оно общее для всех элементов нашего тензора, и область 4×4 из входной картинки (входного тензора) — всего 16 операций чтения. В сумме получаем 25 операций.
MRT на один пиксель. Поскольку все операции чтения — на четыре выходных пикселя, то если пересчитать их на один пиксель, получим чуть больше 6 операций. Это порядка 35% от изначальных 18 операций. Ожидаемый теоретический прирост производительности от использования этой оптимизации — в три раза.
Стоит отдельно пояснить, почему мы используем именно четыре выходных значения. Стандартом гарантировано восемь, но мы решили остановиться на четырёх, потому что при написании такой оптимизации код сильно усложняется. Восемь тоже можно было использовать, но прирост производительности был бы не так ощутим, и мы решили, что оно того не стоит.
Тестируем гипотезу
Протестировать гипотезу решили на реальных слоях нейросети. Мы взяли самые тяжелые слои наших архитектур, чтобы выяснить, что будет, если паковать тензоры в одну текстуру на входе, одну на выходе, и что будет, если использовать четыре текстуры: проработали разные комбинации.
Выяснили, что при упаковке в четыре текстуры, получили существенный прирост производительности — на реальных слоях в два−три раза.
Стоит отметить, что иногда получаются не очень приятные комбинации, в которых, наоборот, происходит замедление. Если такие комбинации нужны вашей архитектуре, за этим нужно следить и вовремя перепаковывать тензоры — это займёт меньше времени, чем возможное замедление.
В результате на замене фона получили 36% реального прироста производительности для нейросети и обработок. Это тот прирост, который увидит конечный пользователь. Это самая эффективная оптимизация из всех, которые мы пробовали.
Как тестировать
Теперь рассмотрим нюансы тестирования таких систем.
Нам нужно было тестировать:
Что нейросети работают правильно (Sanity check).
Что код работает правильно (Unit-тесты).
Производительность.
Sanity check результатов
Мы берём эталонный набор изображений, прогоняем через inference, делаем эталонную реализацию на проверенном движке, например, ONNX Runtime. Получаем два результата и сравниваем их между собой. Если Sanity-check пройден, то наша сеть работает нормально и грубых ошибок в ней нет.
Unit-тесты
Сеть проверили, дальше нужно было запускать юнит-тесты. И тут возникла проблема. Чтобы тестировать приложения, использующие внутри OpenGL, нужно иметь доступную реализацию API OpenGL и устройство дисплея, на котором можно создать контекст. При этом мы хотим эти юнит-тесты запускать в CI-окружении для всех Merge Request. А CI-окружение в нашем случае — это Kubernetes кластер с виртуалками, в которых нет ни дисплея, ни графической карты. Соответственно, ничего из этого запустить мы не могли.
Есть драйвер Mesa на Linux с реализацией API OpenGL, и есть конкретный подвид этого драйвера — Off-Screen Mesa с реализацией на центральном процессоре. Она не очень быстрая, но для unit-тестов подходит.
Для дисплея нам на самом деле нужен не сам экран, а лишь буфер этого экрана, в который можно записать нужные фрагменты для создания контекста. В Linux есть пакет XVFB (X Virtual FrameBuffer), который предоставляет буфер экрана, не требуя его физического наличия. Там мы создаём контекст, пишем в эту память и с ней работаем.
Тестирование производительности
Перед тем как тестировать производительность, вспомните о нескольких неприятных нюансах:
OpenGL — это асинхронный API. Если вы будете запускать замеры вызова функции, то, скорее всего, получите странные цифры. Нужно останавливать замер только тогда, когда вы прочитали из памяти посчитанные данные.
Помимо нагрузки, следите за мощностью. Часто помимо скорости выполнения алгоритма, мы хотим измерять нагрузку: насколько он нагрузит устройство. Здесь тоже есть нюанс. У нас одна из основных платформ — десктоп. В ноутбуках обычно есть центральный процессор, графика и общий бюджет мощности. При этом мы не можем одновременно отдать 100% и центральному процессору, и графике.
Может произойти такая ситуация: например, у нас есть два алгоритма, и тот, и другой показывают 90% нагрузки на графику. Но при этом в первом случае у нас вся доступная мощность уходит на графику, а центральный процессор сидит на базовых частотах и система подтормаживает — это неприятный user experience. Во втором случае графика сидит на своей минимальной мощности, процессор спокойно может разгонять ядра, всё быстро работает, experience классный, а нагрузка и там, и там — 90%.
Поэтому нужно следить за тем, как именно распределяется мощность, как минимум, за тем, как себя ведут частоты на центральном процессоре во время запуска на графике. Иначе можно увидеть не те цифры, что в реальности.
Готовим модели
Напоследок расскажу, как делать модели чуть-чуть более приятными для внедрения в inference.
Есть два разных подхода к работе с моделями:
Исследования моделей:
Берём SOTA-архитектуры из статей.
Часто есть странные решения в архитектуре — сделанные под другие ограничения внешнего мира (бенчмарки и открытые данные).
Если мы — исследователи, то наша задача — сделать модель, которая будет давать хорошие метрики. Для этого берём статьи, смотрим, примеры крутых архитектур — они обычно борются между собой за последние доли процентов в метрике. В таких архитектурах часто есть странные решения, обусловленные конкретным бенчмарком, который пытаются оптимизировать, или некой математической идеей, но с точки зрения организации графа вычислений — не имеющие смысла.
Разработка продукта:
Не используем без надобности нестандартные слои.
Используем проверенные эффективные архитектуры моделей.
При разработке продукта нас интересует другое. Например, как оптимизировать метрику путём работы с данными, аугментациями или другими частями пайплайна. При этом мы не используем без надобности нестандартные решения, стараемся брать проверенные архитектуры, которые лучше скейлятся при работе с данными и которые проще оптимизировать.
Но при подготовке модели всё равно могут возникнуть проблемы. Разберём простой пример, который возникает при исследованиях или попытке внедрить их в продукт.
У нас в продакшене была хорошо работающая оптимизированная модель. Нам хочется её ещё немного ускорить за счёт оптимизации. Мы использовали один из стандартных подходов — прунинг. В результате должны были получить модель поменьше, но с хорошими метриками. Подрезали модель, сделали циклы файн-тюнинга и ожидали, что она будет работать быстрее, не проседая по метрике. Но протестировав, обнаружили снижение скорости inference.
Оказалось, что в полученной после прунинга модели количество каналов перестало быть кратным четырём. Мы порезали каналы, как считали нужным, по метрике, и получили странную архитектуру. Часть оптимизации была завязана на том, что каналы идут кратными четырём, на этом завязана паковка тензоров. Как только количество каналов перестало быть кратным четырём, всё стало работать медленно.
Решение этой проблемы — довольно простое. В прунинг надо добавить правило, что каналы нельзя выкидывать просто так. Это яркий пример, как не думая о том, как архитектура выглядит для конечного исполнения, мы можем на пустом месте потерять перформанс.
Итоги
Важно удерживать всё в памяти графического процессора GPU.
Графовый движок важен, если мы хотим обеспечить оптимальный перформанс.
Кастомный inference на WebGL и некоторые особенности оптимизации дали нам хорошие бусты по производительности и возможность разгрузить центральный процессор для других задач даже при использовании продукта в веб-браузере.