И снова здравствуйте! Продолжаю свой рассказ о поиске «универсального рецепта» для обучения моделей в Luna Line. Теперь речь пойдет о сегментации. Если вы пропустили предыдущую публикацию по классификации, то вам лучше начать читать с нее, чтобы не запутаться. 

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

Новые вводные

Размер датасета: от 10 до 1000 изображений (сегментация размечается еще тяжелее). Мы не привязываемся к семейству. Мы фиксируем три весовые категории (Small, Base, Large) и устраиваем в каждой независимый конкурс: каждая архитектура обучается в своем оптимальном пайплайне, с рекомендованными в оригинальных статьях параметрами и лоссами. Добавляем анализ влияния инициализации (pretrain). Добавляем этап постпроцессинга (морфологические операции), параметры которого подбираются автоматически на валидации.

На картинке суммировала, что у нас поменялось:

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

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

Этап 1 — выбор архитектур. Отдельно подбираем лучшие модели в трех категориях: small, base и large. Для small-моделей дополнительно проверяем влияние предобучения (pretrain).

Этап 2 — подбор гиперпараметров. Для каждой категории (small, base, large) запускаем отдельный поиск. В результате получаем конфиги обучения, адаптированные под размер датасета и конкретную архитектуру.

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

Эксперименты

Этап 1. Конкурс архитектур в трех весовых категориях

Мы собрали пул современных архитектур сегментации: SegFormer, UPerNet, DeepLabV3+, PSPNet и другие. Распределили их по весовым категориям, ориентируясь на число FLOPS, params и inference time. Смотрите Таблицу 1. 

Для small-моделей дополнительно мы проверяли влияние инициализации моделей, а в base категории соревнование по сути не состоялось. 

Результат

Среди Small-моделей лучший результат показала SegFormer-b0. Для base-модели взяли SegFormer-b2, для large — SegFormer-b4.

Как я уже говорила, мы еще тут протестировали влияние претрейна в маленьких модельках — проверили разные стратегии инициализации для SegFormer-b0: случайная инициализация, ImageNet-веса только для энкодера, ADE20K-веса только для энкодера и ADE20K-веса для полной модели (энкодер + декодер).

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

Итог: несмотря на то, что в отдельных точках UPerNet вырывался вперед, мы выбрали SegFormer для всех трех размеров. Причина — универсальность. Нам было важно не плодить сущности и не переключать архитектуру в зависимости от размера модели. SegFormer позволил нам взять B0, B2 и B4 из одного семейства.

Этап 2. Подбор гиперпараметров

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

Здесь мы уже не пытались найти универсальную формулу. Для каждой модели зафиксировали: оптимальный batch size (обычно два, четыре или восемь, упирались в память), Learning rate (индивидуальный), количество эпох (зависит от размера выборки).

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

Этап 3. Постпроцессинг

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

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

Сам постпроцессинг — это базовые морфологические операции (например, opening и closing), которые помогают убирать шум, мелкие артефакты и «зашивать» дырки в масках. За счет этого можно получить дополнительный прирост качества поверх основной модели.

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

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

Итоги по сегментации

В отличие от классификации здесь разрыв между моделями Small (B0), Base (B2) и Large (B4) оказался существенным. Эксперименты показывают, что за счет более тяжелой архитектуры можно выиграть в качестве, особенно когда данных мало.

Все итоги этого блока исследований сведены на рисунке 7.

Вывод по двум этапам исследований

Глобальный вывод, который мы сделали, пройдя оба этапа исследований, звучит банально, но подтвержден тысячами экспериментов: чем меньше данных и чем меньше сама модель, тем более аккуратно нужно подходить к выбору гиперпараметров. На больших объемах данных архитектура и параметры уже не так критичны — данные сами вытянут качество. Но если у нас 10–100 картинок, каждое решение может стоить нескольких процентов точности.

Отсюда главный совет, который мы теперь даем нашим клиентам: не гонитесь за «большими и тяжелыми» моделями. На начальном этапе лучше разметить еще 50–100 изображений, чем пытаться выжать качество из сложной архитектуры. Данные решают.

Ну и еще наше исследование показало, что создание «волшебной кнопки» для обучения моделей — это не просто инженерная задача, а полноценная научная работа. Более 20 000 экспериментов, полгода работы и несколько разочарований (когда любимые архитектуры проигрывали) привели нас к стройной системе. Результат этого: теперь пользователь просто нажимает кнопку «Обучить модель» в Luna Line, а далее система автоматически определяет размер датасета, выбирает модель (RegNetY или ConvNeXt для классификации; SegFormer B0/B2/B4 для сегментации), устанавливает learning rate, batch size и количество эпох по заранее выведенным формулам и опционально по желанию пользователя запускает постпроцессинг. 

Спасибо, что дочитали до конца. Подписывайтесь — постепенно буду делиться результатами исследований по этим и другим задачам компьютерного зрения.