Привет, Хабр!

Хочу рассказать вам нашу историю о том, как изначально рутинная рабочая задача закончилась созданием открытой state-of-the-art нейросети, научной работой и новым датасетом.

С чего всё началось

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

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

И очень нехорошо, если на фото ребёнок, а мы предложили купить взрослую одежду, или, например, перепутали пол. А по самой одежде это очень тяжело понять (хотя мы и так тоже пытаемся, но точность невелика). Кроме того, возраст может в принципе давать ещё какую-то дополнительную информацию, помогающую делать рекомендации качественнее.

Вообще, open source решений для этой задачи довольно много. До создания своего, мы использовали открытую модель FairFace, выбрав её за простоту и близость к нашей задаче. Но, процент ошибок в проде удручал - в разных проектах он составлял от 1 до 10% и более. Соотвественно, это потенциальные 10% случаев, когда рекомендация вообще никак не вкатит пользователю и прибыль будет упущена. Звучит плохо. С другой стороны, мы никогда не занимались лицами профессионально, поэтому старались поступать мудро и надеялись на чужие решения (внутри компании или внешние open source).

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

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

Ход работ: Создание baseline

Решение от FairFace имеет ряд серьёзных недостатков. Наибольшим изначально виделся их классификационный подход. Во-первых диапазоны разбиты неудачно для нас, во-вторых постоянно всплывали проблемы пограничных случаев, в-третьих, алгоритм не делает различия в ошибке между классами, например, 0-5 лет или 60-70, при тренировке, что никак не могло сказаться благоприятно на способности к генерализации. Тренировочного кода к модели тоже нет.

Задача для baseline была такая: небольшая, шустрая модель, работающая с лицами и предсказывающая в один проход сразу пол и возраст. Поэтому, мы начали с timm, огромного репозитория классификационных моделей от Hugging Face, претренированных на больших открытых датасетах.

Мы заменили классификацию на регрессию, с соответствующими изменениями в нужных частях кода. Этого нет в статье, поскольку она посвящена уже конечному решению на трансформере, но на этом этапе нашей основной моделью для экспериментов была чисто свёрточная нейросеть resnext50_32x4d. Она быстрая и хорошо себя зарекомендовала. Кроме того, мы прикрутили к сети LDS и FDS из статьи Deep Imbalance Regression. Не вдаваясь в детали, эти подходы позволяют существенно компенсировать естественный дисбаланс возраста в данных. Особенно полезной оказалась первая техника, позволяющая посчитать веса для примеров в соответствии с их распределением, которые затем применяют в целевой функции MSE (Mean Squared Error).

Первые же эксперименты показали перспективность подхода, например, мы с полутыка получили MAE около 5.0 на IMDB-clean, где SOTA была 4.68.

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

Мы развернулись в сторону трансформеров, поскольку эти модели являются более общими и перспективными, и попробовали разные архитектуры. Многие, например CaiT и XCiT, показали себя очень хорошо, решив проблему недостающей мощности. Но, с ними есть уже другая проблема - они не блещут скоростью. А самый первый (исторически из удачных), весьма быстрый и популярный визуальный трансформер ViT не очень бодро сходился. Скорее всего из-за недостатка данных, этот трансформер очень прожорлив и требует огромных объёмов.

Через какое-то количество экспериментов, мы наконец нащупали идеальный вариант. Им стала гибридная модель VOLO. Она сочетает в себе преимущества свёрточных и трансформерных нейросетей. В этой модели, вместо простого нарезания изображения на патчи применяется сначала ряд свёрточных слоёв. Кроме того, в ней используется особый блок внимания - Outlook Attention, который помогает решить часть проблем, возникающих у этих архитектур при адаптации к изображениям. Эта одна из самых быстрых моделей для визуальных трансформеров, как в плане скорости работы, так и сходимости обучения. И, главное, глохнуть она совершенно не собиралась. Более того, добавив в выход пол, мы поймали желаемый эффект и получили прирост в точности возраста!

Сбор данных

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

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

Поэтому, мы отправились размечать с нуля на краудсорс ресурс Яндекс Толока. У читателя, вероятно, возникнет вопрос, как же мы собрались это делать, если сами только что упомянули невысокую точность человеческой разметки? Всё так.

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

А поскольку любой краудсорс регулярно подвергается набегам ботов и читеров, плюс, чтобы потом оценить человеческую точность и найти лучший метод агрегации голосов, мы положили в каждый набор из 6-ти примеров 7-ой, контрольный. Эти контрольные задания мы набрали из IMDB-clean датасета и знали для них точные ответы. А вот исполнители, напротив, не знали об этом ничего.

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

Суммарно мы собрали около 500 000 изображений из нашего прома и из Open Images Dataset. Часть данных из последнего было решено выложить, чтобы у сообщества исследователей наконец появился по-настоящему сбалансированный регрессионный бенчмарк. Мы сбалансировали его не только в общем, но и по полу внутри 5-летних диапазонов возрастов:

Что же касается агрегации голосов, то в итоге были перепробованы почти все существующие методы:

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

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

MiVOLO: успех выходит из под контроля

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

И тут встал вопрос, как учить сеть делать две задачи (ещё и с мульти-таск выходом) и не растерять всю точность? А желательно бы ещё её и улучшить.

Размер входа у сети весьма скромный - 224х224. Если просто подавать туда весь образ человека - тело с лицом, то нейросеть гарантированно круто просядет в работе по лицам, а это самый надёжный вариант. Увеличивать разрешение? Очень дорого, скорость упадёт катастрофически.

Значит, правильно было бы подавать эти изображения независимо, как два входа в два свёрточных стебля (conv-stem), а потом объединять.

Большой вопрос как и когда это делать. Если использовать late fusion и собирать признаки где-то в конце сети, потеряем или в скорости (придётся делать два параллельных бранча) или в точности (сделав два бранча, но порезав их размеры и параметры), а заодно не сможем использовать transfer learning и претренированные веса из-за изменения размерностей. Следовательно, точно нужен early fusion, с понижением размерностей до исходных.

Эксперименты мы начали с классики, которая давно применяется в свёрточных сетях - 1х1 конволюционного squeeze слоя, который получает на вход 2N каналов, а отдаёт N. Этот вариант работал, но был хуже в точности на тех примерах, где лица есть. Казалось, что можно лучше. Поэтому, мы перепробовали множество вариантов, например, неплохим был BottleneckAttention, но, в итоге, спроектировали свой собственный модуль.

На изображении можно увидеть основную суть, как мы это решаем. Мы берём сжатые представления от лиц и тел (по сути это уже патчи), выполняем cross-attention сначала в одну сторону, затем в другую, после чего признаки объединяем и сжимаем количество каналов через MLP:

Таким образом, модуль решает три проблемы проблемы разом:

1. С помощью механизма внимания признаки обогощаются и становятся качественнее

2. Достигается целевая размерность признаков

3. Эффективно обрабатываются примеры, когда полезная информация на изображении занимает не всю его часть. И те случаи, когда вход и вовсе пустой (см. ниже). У чистых трансформеров с такими примерами большие проблемы.

Решение сразу показало себя очень хорошо, и наконец принесло желаемое - сеть давала во всех случаях результаты лучше, чем VOLO с одним входом.

Оставалась последняя проблема - тренируется наша модель очень долго. Если VOLO требовал около 300 эпох, чтобы сойтись, то MiVOLO уже 700. На полном датасете из полумиллиона изображений это было весомое время даже на мощном сервере NVIDIA DGX. Поэтому, мы немного упростили задачу для сети и стартовали не из Imagenet чекпоинта, а сразу из нашего же VOLO, натренированного как baseline. Поскольку веса есть только для одного из стеблей, мы просто продублировали их. А заодно заморозили тот, который для лиц - он уже достаточно хорош и нет смысла ничего в нём менять. Этот вариант сходится за дополнительные за 250-300 эпох, с пониженным learning rate. Итоговые вычислительные затраты не так велики, и часть экспериментов мы выполняли на простеньком сервере с двумя NVIDIA A4000.

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

Объединив всё это и все собранные данные, получилось так хорошо, что пробили и без того SOTA результат на IMDB, а на UTKFace взяли первое место без доп. данных и почти без изменений в тренировке. Больше регрессионных датасетов, которые были бы солидных размеров и содержали пол, не оставалось.

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

Но, нас ждал сюрприз. Ничего не меняя, просто запустив готовую сеть, мы сразу взяли первое место и на одном из самых популярных и старых датасетов Adience и на свежем FairFace. Больше побеждать особо и некого - остальные датасеты или тяжело достать (нужно заполнять формы, авторы не очень уже и отвечают и пр.) или слишком маленькие и не интересные.

Второй сюрприз заключался в точности работы системы на примерах без лиц. Для замера, мы закрыли чёрным квадратом лица в тестовых примерах и были сильно удивлены. Получился MAE 6.66, что лучше, чем человеческая точность с видимым лицом (см. ниже)!

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

Но итоговая точность не тянет на "хотя бы как-то", она скорее "огого":

Это говорит о том, что модель научилась генерализировать по-настоящему, её понимание задачи выходит за рамки обучающей выборки.

Все итоговые результаты и бенчмарки можно посмотреть на Papers With Code.

Кстати, я мало тут говорю про точность пола, т.к. это чуть менее интересная и чуть менее сложная задача, чем возраст, но с ним мы так же взяли первое место на Adience, причём, с совсем уже каким-то космическим отрывом:

Точность человека в задаче возраста

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

Получается, что средняя точность 7.22, медиана отличается не сильно из-за формы распределения, близкой к симметричной. Это очень большая ошибка.

Кстати, лучший индивидуальный результат в левом хвосте - 4.54.

Интересно, а как оценивают свою точность люди? Я провёл опрос в своих соц. сетях и собрал 105 голосов, попросив оценить, какую среднюю ошибку респонденты ожидают у людей в такой задаче. Вот результат:

Так что самокалибровка у нас довольно неплохая.

Самый важный вопрос - насколько модель точнее человека? Намного:

А самый популярный вопрос, который мне задают про точность людей, это есть ли связь с возрастом разметчика? Это интересный момент. Но оказалось, что Толока не отдаёт эту информацию по API. При этом, в веб версии она есть. Ну, чтож, пришлось расчехлять Selenium.

Результат:

Как видно из графиков, связи практически нет. Для разметчиков старше 60 - 65 лет я не стал строить график - слишком мало данных и результаты там "скачут". Единственное, что обращает на себя внимание, это заметно пониженная точность в диапазоне 20 лет. Что это: юношеская нетерпеливость при разметке или особенности восприятия людей старше себя в этом возрасте? Я оставляю предполагать читателю. Особенно с учётом, что пользователи сервиса указывают возраст самостоятельно и его никто не проверяет.

Ещё, интересным может быть такой вопрос - улучшалась ли точность людей в процессе разметки? Поскольку никакой обратной связи, кроме перманентного бана, не было, крайне маловероятно. Но всё-таки я проверил, вдруг у нас в мозге в таких случаях работает какой-то механизм unsupervised обучения. Однако, чуда не произошло. И даже наоборот - точность, в отсутствие каких-то санкций за небрежную разметку, падала (и в итоге разметчик попадал в бан). Если из этого и следует какой-то вывод, так тот, что в идеале очень желательно давать источник обратной связи. С другой стороны, это довольно трудозатратно реализовать и это всегда вопрос желаемых затрат и профита. В нашем случае овчинка выделки не стоила.

Послесловие

Есть много вещей, которые могут серьёзно улучшить модель. Главное, что я бы подчеркнул - это работа с кропом тела. Использованный нами метод достаточно топорный, основанный на простой обрезке через классический Image Processing. Если прикрутить к пайплайну вышедший на днях Fast SAM, позволяющий получать точные маски объектов, то можно добиться серьёзного улучшения. Кроме того, остаются проблемы с предсказанием возраста после ~70, но это решается через данные.

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

Та же история с некоторыми фильтрами, например, многие эффекты в ТикТок, да даже просто хорошее сглаживание, результат сильно изменят.

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

Материалы

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

Материалы из этого текста:

• Сама статья на Arxiv

• Репозиторий с кодом для валидации и инференса.

• Демо на Hugging Face

• Лендинг нового датасета и дополнительных данных.