Swin Transformer - 이론과 구현

Swin Transformer 개요

Swin Transformer란 2021년에 발표한 대규모 비전 인식 모델이다. 해당 아키텍쳐가 제안된 논문은 "Swin Transformer: Hierarchical Vision Transformer using Shifted Windows"이다. 기존 트랜스포머를 기반으로 하는 모델의 공통적인 문제점은 고정된 패치 크기로 인해 세밀한 예측을 필요로 하는 Semantic Segmentation과 같은 작업에 적합하지 않다는 점이었다. 또한 트랜스포머의 Self Attention은 입력 이미지 크기에 대한 Quadratic 계산 복잡도를 가지므로 고해상도 이미지를 처리하기 어려웠다. 이러한 문제를 극복하기 위해 Hierarchy 특징 맵을 구성해 Linear 계산 복잡도를 갖는 Swin Transformer가 등장했다.

Swin Transformer의 대표적인 장점은 다음과 같다:

• 지역(window) self-attention + shift 연산을 활용해 계산량을 줄이고 성능을 높인 구조
• 계층마다 어텐션이 수행되는 패치의 크기와 개수를 계층적으로 적용해 처리하므로 높은 해상도와 다양한 크기를 가진 객체를 효율적으로 처리
• 객체 탐지나 객체 분할 등과 같은 작업에서 백본 모델로 사용

ViT 와 Swin Transformer의 차이

이미지 패치 비교 Swin Transformer, ViT https://arxiv.org/pdf/2103.14030

ViT 모델의 한계점

• ViT 모델은 획일적인 패치 크기와 위치에 대한 제약으로 다양한 이미지 크기와 종횡비를 가진 데이터 세트에 대해서는 적용이 어렵다.
• 또한 패치 단위로 입력을 받기 때문에 이미지의 공간 정보가 완전히 보존되지 않을 수 있다. 
• 이미지의 패치 수가 증가하면 학습데이터의 크기도 커지기 때문에 모델 학습에 필요한 계산량이 증가할 수 있다.
• 대규모 데이터세트에서는 학습에 필요한 컴퓨팅 자원을 많이 사용하게 되는 단점이 있다. 

Swin Transformer의 특징

• 스윈 트랜스포머는 로컬 윈도(Local Window)를 활용해 물체의 크기나 해상도를 계층적으로 학습, 로컬 윈도는 스윈 트랜스포머에서 입력 이미지를 처리하기 위해 사용되는 고정 크기의 작은 윈도우
• 스윈 트랜스포머는 이러한 로컬윈도를 이동시키는 시프트 윈도 기술을 도입해 입력이미지를 일정한 크기의 패치를 분할하고, 특징 맵을 이동시켜 다음계층에 사용함으로서 패치의 위치를 더 유연하게 다룸 -> 이를 통해 더 높은 정확도와 더 적은 매개변수 수를 가지면서 더 넓은 범위의 이미지 크기와 종횡비를 다룰 수 있음
• 계층적 접근 방식은 로컬 윈도가 각 패치의 세부 정보에 집중하게 되고 시프트 윈도우를 통해 전역 특징을 확인할 수 있으므로 고해상도 이미지를 효율적으로 처리 

Swin Transformer 모델 구조

Swin Transformer는 이미지를 패치 파티션으로 구분한 후, 선형 임베딩(Linear Embeddinng), SwinTransformerBlock, 패치 병합 연산이 반복적으로 수행된다. 

스윈 트랜스포머 구조 https://arxiv.org/pdf/2103.14030

 

• 먼저 입력 이미지를 일정한 크기의 패치로 분할하고, 각 패치에 대해 선형 임베딩을 수행한다. 이를 통해 각 패치는 고정된 차원의 벡터로 변환된다. 다음으로 분할된 패치들을 기반으로 스윈 트랜스포머 블록을 구성한다. 이 블록은 Attention, Layer Normalization, MLP를 포함하며 패치간 상호작용을 수행한다.
• 이후 이 패치를 병합해 전체 이미지에 대한 분류를 수행한다.  SwinTransformer는 일반적으로 네개의 스테이지로 구성되어 각 스테이지는 서로다른 해상도와 패치 크기를 갖는다. 앞서 설명한 선형 임베딩, 스윈 트랜스 포머 블록, 패치 병합 등이 반복해 적용된다.
• 스테이지 1에서는 패치들이 선형 임베딩된 후 스윈 트랜스포머 블록에 전달되며, 스테이지 2,3,4에서는 패치 병합 후 각 스윈 트랜스포머 블록에 전달된다. 패치 병합은 모델의 매개변수를 줄이기 위해 이미지 텐서 C,H,W를 2C,H/2,W/2로 재정렬해 저차원 임베딩을 수행한다.
• 위 그림에서 Patch Partition이란 입력 이미지를 작은 사각형 패치로 분할해 처리하는 방식을 의미한다. 분할된 패치는 트랜스포머 계층의 입력으로 사용되며, 입력이미지의 공간 정보를 보존하고 트랜스포머 계층에 대한 계산효율성을 높인다. 코드로 구현하면 다음과 같다:

import torch

def patch_partition(x:torch.Tensor, window_size:int):
    '''
    Args:
    x: (B,H,W,C) - 입력 feature Map
    window_size: 윈도우 크기
    return:
        windows : (num_windows * B, window_size, window_size,C)
    '''

    B,H,W,C = x.shape

    # H,W 차원을 window 단위로 reshape
    x = x.view(
        B,
        H // window_size, window_size,
        W // window_size, window_size,
        C
    )

    # 윈도우를 모양에 맞게 
    # congiguous : Tensor의 각 값들이 메모리에도 순차적으로 저장되어 있는지 여부
    windows = x.permute(0,1,3,2,4,5).contiguous().view(
        -1, window_size,window_size,C
    )
    return windows

def window_reverse(windows, window_size, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, C)
        window_size: 윈도우 크기
        H, W: 원래 height, width
    Returns:
        x: (B, H, W, C)
    """
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.view(
        B,
        H // window_size, W // window_size,
        window_size, window_size, -1
    )

    # 윈도우를 원래 이미지 순서로 재배열
    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    return x

• window_partition()을 통해 [B,H,W,C]에서 [num_windows*B, window_size, window_size, C] 형태로 나눠서 윈도우 단위 Self Attention을 적용하게 된다. 이후 MLP 혹은 Residual Connection 등 전체 이미지 기반 연산을 위해 다시 원래 형태로 복원해야하는데 이때 Window_reverse()를 사용한다.
• SwinTransformer Block은 패치 파티션 이후 네 개의 스테이지 안에서 반복 적용되며, 선형 임베딩 또는 패치 병합 이후에 수행된다. 스윈 트랜스포머 블록은 계층 정규화, Window Multi Head Self Attention, MLP, Shifted Window Multi Head Self Attention 등으로 구성된다.

Swin Transformer Block 구조

• 스윈 트랜스포머 블록 안에서는 트랜스포머 모델에서 사용되던 MSA(Multi Head Self Attention) 방식이 아닌 W-MSA와 SW-MSA로 구성된다
• W-MSA란 로컬 윈도를 사용하는 멀티 헤드 셀프 어텐션으로 입력 특징 맵을 중첩되지 않게 나눈 다음 각 윈도에서 독립적으로 셀프 어텐션을 수행한다. W-MSA는 이미지에서 일정한 크기의 윈도우 영역을 설정하고, 해당 영역 내 픽셀 간의 어텐션을 계산한다. 이를 통해 이미지 내의 전체적인 어텐션을 수행하고 특정 위치에 대한 정보를 추출한다. 또한 각 윈도우 내의 지역적인 특징을 분석할 수 있으며 전체 입력에 대한 셀프 어텐션 비용을 줄인다.
• W-MSA 방식은 연산량을 대폭 감소 시켰지만 윈도우 내부의 이미지 패치 영역에서만 셀프 어텐션을 수행한다는 한계가 있어서 SW-MSA는 윈도우를 가로 또는 세로 방향으로 이동한 후, 인접한 윈도우 간의 셀프어텐션을 계산한다.
• SW-MSA는 로컬 윈도우 간의 셀프 어텐션을 수행하기 위해 W-MSA에서 사용된 로컬 윈도우 개수가 많아서 더 많은 연산을 요구한다. 이 문제를 해결하기위해 순환시프트와 어텐션 마스크를 사용했다. 다음은 그 그림이다.

효율적인 배치 계산, 순환 시프트

• 순환 시프트는 로컬 윈도우가 이동할 때 이동한 위치의 정보를 이전 위치에서 가져오는 것을 말하며 로컬 윈도우 사이즈 M보다 작은 M/2 만큼 로컬 윈도우들을 이동시킨다. SwinTransformerBlock을 Pytorch로 구현하면 다음과 같다

class SwinTransformerBlock(nn.Module):
    def __init__(self,dim,input_resolution,num_heads,
                 window_size=7,shift_size=0,mlp_ratio=4.,drop_out=0.1):
        super(SwinTransformerBlock, self).__init__()
        self.dim = dim # 입력/출력 채널 수
        self.input_resolution = input_resolution # 입력 Feature map의 해상도
        self.window_size = window_size # 윈도우 크기
        self.shift_size = shift_size # Shifted Window 일 경우 이동 크기 (보통 window_size//2)
        self.num_heads = num_heads

        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(
            dim,num_heads,dropout=drop_out,batch_first=True
        )
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim,int(dim*mlp_ratio)),
            nn.GELU(),
            nn.Dropout(drop_out),
            nn.Linear(int(dim*mlp_ratio),dim),
            nn.Dropout(drop_out)
        )

    def forward(self,x:torch.Tensor):
        '''
        x : (B, H*W,C)
        shift_size = 0: W-MSA
        shift_size > 0: SW-MSA
        '''
        H,W = self.input_resolution
        B, L, C = x.shape
        assert L == H * W, "input feature has wrong size"

        x = x.view(B, H, W, C)
        
        if self.shift_size > 0:
            # X를 오른쪽 아래로 이동
            #  torch.roll 은 원하는 dim으로 tensor를 이동
            shifted_x = torch.roll(x,shifts=(-self.shift_size,-self.shift_size),dims=(1,2))
        else:
            shifted_x = x

        # 윈도우 단위로 나누기
        x_windows = patch_partition(shifted_x,self.window_size)
        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)

        # 윈도우 단위로 MultiheadAttention 수행
        shortcut = x_windows
        x_windows = self.norm1(x_windows)
        attn_windows ,_ = self.attn(x_windows,x_windows,x_windows)
        x_windows = shortcut + attn_windows # residual connection
        
        # 윈도우 병합
        x_windows = x_windows.view(-1, self.window_size , self.window_size, C)
        x_windows = window_reverse(x_windows, self.window_size, H, W)

        # x를 원래 위치로 되돌리기
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x

        x = x.view(B, H * W, C)

        # MLP 수행
        x = x + self.mlp(self.norm2(x))
        return x

• 위의 블록을 번갈아 쌓으면 Swin 구조가 완성된다.
• 정리하자면 SwinTransformer는 위에서 설명한 구조 대로 지역 Self-Attention + Shift 연산을 활용해 계산량을 줄이고 성능을 높인 구조이다.