Реализация SegFormer в PyTorch

Реализация SegFormer в PyTorch

Быстрая, эффективная и легкая модель для сегментации изображений

Привет!! Сегодня мы увидим, как реализовать SegFormer в PyTorch, предложенном в статье SegFormer: простой и эффективный дизайн для семантической сегментации с преобразователями.

Код находится здесь, интерактивную версию этой статьи можно скачать здесь.

Давайте начнем!

В статье предлагается новая модель на основе преобразователя для сегментации изображений. Даже если слово «трансформер» сейчас модное, а сама модель имеет только базовый механизм внимания. Эта модель имеет два основных преимущества: во-первых, SegFormer содержит новый кодировщик Transformer с иерархической структурой, который выводит многомасштабные функции. Затем не требуется позиционное кодирование, что позволяет избежать интерполяции позиционных кодов, которая приводит к снижению производительности, когда разрешение тестирования отличается от разрешения обучения.

Как ни странно, мы возвращаемся в исследованиях назад, эти два преимущества присутствуют в коннетах с самого начала, и мы увидим, что SegFormer, в конце концов, представляет собой просто коннет + внимание.

На следующем рисунке показана производительность SegFormer по сравнению с различными моделями/размерами в наборе данных ADE20K, у них есть sota.

Он лучше старого доброго FCN-R50 и в 2 раза быстрее. Поскольку у него на 24 FLOPS меньше, мне интересно, почему он только в два раза быстрее.

Архитектура

Модель представляет собой классический кодер-декодер/магистраль-шея. Голова прикреплена для предсказания окончательной маски сегментации.

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

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

Вот как это должно выглядеть:

С небольшим фотошопом

Посмотреть код здесь

Декодер

Используемый декодер называется MixVisionTransformer (MiT), другой ViT с каким-то случайным материалом посередине, мы назовем его SegFormerDecoder. Начнем с первого отдельного компонента самого блока, OverlapPatchMerging.

OverlapPatchMerging

Блок OverlapPatchMerging может быть реализован со слоем свертки с stride меньшим, чем kernel_size, поэтому он перекрывает разные патчи. Это то же самое, что было предложено много лет назад, когда вы используете stride больше единицы, чтобы уменьшить пространственное измерение ввода. В SegFormer за конверсионным слоем следует норма слоя.

Поскольку nn.LayerNorm в PyTorch работает для тензоров формы batch, ...., channels, мы можем создать LayerNorm2d, который сначала меняет местами ось channels с последней, затем применяет норму слоя и меняет ее обратно. Я буду использовать einops, чтобы сделать код более читабельным.

Тогда наш OverlapPatchMerging — это просто слой conv, за которым следует норма нашего слоя.

Эффективное внимание к себе

Мы все знаем, что внимание имеет квадратную сложность O(N^2), где N=H*W в нашем случае. Мы можем уменьшить N в R раз, сложность становится O(N^2/R). Один из простых способов — сгладить пространственное измерение и использовать линейный слой.

Мы уменьшили пространственный размер на r=4, то есть на 2 в каждом измерении (height и width). Если вы думаете об этом, вы можете использовать слой свертки с kernel_size=r и stride=r для достижения того же эффекта.

Поскольку внимание равно softmax((QK^T/scale)V), нам нужно вычислить K и V, используя редуцированный тензор, иначе формы не будут совпадать. Q \in NxC, K \in (N/R)xC, V \in (N/R)xC, мы можем использовать MultiheadAttention PyTorch для вычисления внимания.

torch.Size([1, 8, 64, 64])

МиксМЛП

Внимательный читатель мог заметить, что мы не используем позиционное кодирование. SegFormer использует 3x3 глубинную конверсию. Цитата из статьи Мы утверждаем, что позиционное кодирование не обязательно для семантической сегментации. Вместо этого мы представляем Mix-FFN, который учитывает влияние заполнения нулями на утечку информации о местоположении. Я понятия не имею, что это значит, поэтому будем считать это само собой разумеющимся.

Я почти уверен, что он называется Mix, потому что он смешивает информацию, используя 3x3 conv.

Слой состоит из dense layer -> 3x3 depth-wise conv -> GELU -> dense layer. Как и в ViT, это обратный слой узкого места, информация расширяется в среднем слое.

Блок энкодера (преобразователя)

Давайте соберем все вместе и создадим наш блок кодировщика. Мы будем следовать лучшему (имхо) соглашению об именах, мы называем SegFormerEncoderBlock часть с вниманием к себе и миксом-fpn и SegFormerEncoderStage слияние всего патча с перекрытием + N x SegFormerEncoderBlock

Очень похоже на ViT, у нас есть пропущенные соединения и слои нормализации + Stochastic Depth, также известный как Drop Path (у меня есть статья об этом).

torch.Size([1, 8, 64, 64])

Хорошо, давайте создадим сцену. Не знаю почему, они в конце применяют норму слоя, так что будем делать так же :)

Окончательный SegFormerEncoder состоит из нескольких этапов.

Я добавил функцию chunks, чтобы код оставался чистым. Это работает так

[[1, 2], [3, 4, 5]]

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

Из кодировщика мы возвращаем список внутренних признаков, по одному на каждом этапе.

Декодер / Шея

К счастью, изображение декодера/грифа совпадает с исходным кодом. Они назвали часть декодера MLP Layer.

То, что он делает, очень просто: он берет F объектов с размерами batch, channels_i, height_i, width_i и выводит F' объектов с одинаковым пространственным и канальным размером. Пространственный размер фиксируется на first_features_spatial_size / 4. В нашем случае, поскольку наш вход представляет собой изображение 224x224, на выходе будет маска 56x56.

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

В данном случае у нас нет этапов, поэтому наш SegFormerDecoder — это просто список блоков. Он принимает список объектов и возвращает список новых объектов с тем же пространственным размером и каналами.

Глава SegFormer

Мы почти на месте! Функции декодера объединены (помните, что все они имеют одинаковые каналы и пространственные измерения) на оси канала. Затем они передаются в головку сегментации, чтобы уменьшить их с channels * number of features до channels. Наконец, плотный слой выводит окончательную сегментацию.

SegFormer

Ну, наша окончательная модель просто encoder + decoder + head. Очень просто

Давай попробуем!

torch.Size([1, 100, 56, 56])

Вывод правильный, мы ожидаем маску пространственной формы image_size // 4 и 224 // 4 = 56.

Мы сделали это! 🎉🎉🎉

Выводы

В этой статье мы шаг за шагом рассмотрели, как создать SegFormer; быстрая и эффективная модель сегментации изображений.

Спасибо, что прочитали это!

Франческо