Patch Merging of SWIN

Patch Merging

Patch Merging Module (or Layer) 는

• CNN에서 널리 사용되는 Max (or Averaging) Pooling과 유사한 downsampling 역할을 수행하되,
• 최댓값(or 평균값)을 선택하는 대신
• $2 \times 2$ 인접 patch의 feature를 결합하고 선형 변환하여
• 공간 해상도는 감소시키고 채널 수는 증가시킴으로써,
• Swin Transformer가 CNN과 유사한 hierarchical representation을 학습하도록 함.

---

Patch Merging 의 위치

첫번째 stage를 제외하고 이 후 stage 의 앞에 위치하며, 해상도를 가로세로로 1/2씩 줄이고, feature의 차원을 2배로 올림.

Stage 내에서는 dimension이 유지되고 Stage가 달라지면 Patch Merging에 의해 dimension이 달라짐.

다음은 TorchVision 에서 Swin Base에서의 dimension의 변화를 보여줌(위의 그림은 Swin Tiny)

입력
(B, 3, 224, 224)

Conv2d patch embedding
(B, 128, 56, 56)

Permute([0, 2, 3, 1])
(B, 56, 56, 128)

Stage 1
(B, 56, 56, 128)

Patch Merging
(B, 28, 28, 256)

Stage 2
(B, 28, 28, 256)

Patch Merging
(B, 14, 14, 512)

Stage 3
(B, 14, 14, 512)

Patch Merging
(B, 7, 7, 1024)

Stage 4
(B, 7, 7, 1024)

더보기

 

Swin T(iny) 와 Swin B(ase)의 차이는 다음을 참고:

항목	Swin-T	Swin-B
논문 발표	2021	2021
Patch size	$4 \times 4$	$4 \times 4$
Window size	$7 \times 7$	$7 \times 7$
Stage 수	4	4
Block 수	(2, 2, 6, 2)	(2, 2, 18, 2)
Embedding dimension	96	128
Stage별 채널 수	96 → 192 → 384 → 768	128 → 256 → 512 → 1024
Attention head 수	(3, 6, 12, 24)	(4, 8, 16, 32)
파라미터 수	약 28M	약 88M
FLOPs (224×224)	약 4.5G	약 15.4G

 

---

Patch Merging 의 동작

참고로, Swin V2 에선 조금 차이가 있는데,

Layer Normalization이 Linear Layer를 통과하고 나서 가해짐.

 

다음은 timm 에서의 구현 코드임:

class PatchMerging(nn.Module):
    def __init__(self, input_resolution, dim, out_dim=None, norm_layer=nn.LayerNorm):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.out_dim = out_dim or 2 * dim
        self.norm = norm_layer(4 * dim)
        self.reduction = nn.Linear(4 * dim, self.out_dim, bias=False)

    def forward(self, x):
        """
        x: B, H*W, C
        B: Batch size 
        """
        H, W = self.input_resolution
        B, L, C = x.shape
        x = x.view(B, H, W, C)

        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C

        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C      
        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C

        x = self.norm(x)
        x = self.reduction(x)
        return x

---