Swin Transformer

Swin Transformer (сокращение от Shifted Window Transformer, трансформер со сдвигаемыми окнами) — иерархический визуальный трансформер, который вычисляет самовнимание внутри небольших локальных окон и сдвигает их между соседними слоями. Ze Liu и коллеги из Microsoft Research Asia представили его в статье 2021 года «Swin Transformer: иерархический визуальный трансформер со сдвигаемыми окнами», получившей премию Marr за лучшую статью на ICCV 2021. Благодаря локальным окнам стоимость внимания растёт линейно с числом пикселей, а поэтапное объединение патчей создаёт многоуровневые карты признаков, подобные картам свёрточной нейронной сети. Эти особенности позволили использовать трансформеры в качестве универсальных бэкбонов для детекции объектов и сегментации изображений высокого разрешения, а не только для классификации.

Принцип работы сдвигаемых окон#

Сначала Swin делит изображение на патчи размером 4×4 пикселя и проецирует каждый из них в эмбеддинг токена, поэтому изображение размером 224×224 превращается в сетку токенов 56×56 с 96 каналами в модели Tiny. Глобальное внимание для всех 3 136 токенов сопоставляло бы каждый токен со всеми остальными; Swin разбивает сетку на неперекрывающиеся окна размером 7×7 токенов и вычисляет внимание только внутри каждого окна. В статье стоимость глобального внимания для сетки h×w с C каналами выражена как 4hwC² + 2(hw)²C, а для окон размером M×M токенов — как 4hwC² + 2M²hwC; последняя формула линейна по hw.

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

Между этапами слой слияния патчей конкатенирует каждую группу из 2×2 соседних токенов и линейно проецирует результат, уменьшая разрешение вдвое и удваивая число каналов. Четыре этапа выдают карты признаков с разрешением 1/4, 1/8, 1/16 и 1/32 от входного — с теми же шагами, что у ResNet. Поэтому Swin можно подключить к уже существующим головам детекции и сегментации, например к пирамидальной сети признаков. В реализации torchvision показаны оба механизма:

import torch
from torchvision.models import swin_t

# Swin-T architecture; random weights are enough to inspect its structure
model = swin_t().eval()

# Each pair of blocks uses 7x7 windows, the second shifted by 3 patches
regular, shifted = model.features[1]
print(f"Window {regular.attn.window_size}, shift {regular.attn.shift_size} then {shifted.attn.shift_size}")

# features alternates stages with patch merging; stage outputs are (batch, height, width, channels)
x = torch.randn(1, 3, 224, 224)
with torch.no_grad():
    for i, layer in enumerate(model.features):
        x = layer(x)
        if i % 2 == 1:
            print(f"Stage {i // 2 + 1}: {tuple(x.shape)}")

Выводится Window [7, 7], shift [0, 0] then [3, 3], а затем выходы этапов: (1, 56, 56, 96), (1, 28, 28, 192), (1, 14, 14, 384) и (1, 7, 7, 768). На каждом этапе размер сетки уменьшается вдвое, а число каналов удваивается, формируя пирамиду, которая нужна головам плотного предсказания.

Swin Transformer, ViT и CNN#

Исходный ViT делит изображения на патчи 16×16, сохраняет одинаковое разрешение во всех слоях и использует глобальное внимание. Такой подход подходит для классификации изображений, но становится затратным при высоких разрешениях, необходимых для детекции и сегментации. Swin возвращает два априорных свойства свёрточных сетей — локальность и иерархию признаков, — сохраняя при этом зависящие от входных данных веса механизма внимания. Его локальные окна отличаются от внимания со скользящим окном, при котором каждый токен обращается к своему соседству: в Swin используются фиксированные разделы, общие для токенов, что эффективно отображается на умножения матриц на GPU.

На момент выпуска крупнейшая модель Swin-L, предварительно обученная на ImageNet-22K, достигла точности top-1 87.3% на ImageNet-1K, AP рамок 58.7 и AP масок 51.1 на COCO test-dev, а также mIoU 53.5 на ADE20K — лучших опубликованных результатов в этих тестах детекции и семантической сегментации. Успех модели также дал новый импульс исследованиям свёрточных сетей: ConvNeXt от Meta (2022) модернизировала ResNet, позаимствовав проектные решения у Swin, и сравнялась с ней по точности или превзошла её.

Swin можно использовать в качестве бэкбона пользовательской модели Ultralytics: модуль TorchVision в YAML-файле модели загружает swin_t из torchvision с предобученными весами ImageNet, а голова детекции YOLO обучается на единственной возвращаемой карте признаков с шагом 32.

Варианты и области применения#

• Swin Transformer V2: В 2022 году архитектуру масштабировали до 3 миллиардов параметров и изображений размером до 1 536×1 536 пикселей, используя косинусное внимание с масштабированием и непрерывное смещение позиции с логарифмической шкалой для стабильного обучения.
• Распознавание действий в видео: Video Swin Transformer распространяет окна на три измерения — пространство и время — для классификации действий в видеоклипах.
• Сверхразрешение и восстановление: SwinIR использует блоки Swin для увеличения изображений, удаления шума и артефактов JPEG.
• Анализ медицинских изображений: Swin-Unet и Swin UNETR объединяют энкодеры Swin с декодерами в стиле U-Net; Swin UNETR сегментирует опухоли мозга в трёхмерных объёмах МРТ.