Swin Transformer 코드리뷰

# 인코더 레이어 구성 self.layers = nn.ModuleList() for i_layer in range(self.num_layers): layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer), input_resolution=(patches_resolution[0] // (2 ** i_layer), patches_resolution[1] // (2 ** i_layer)), depth=depths[i_layer], num_heads=num_heads[i_layer], window_size=window_size, mlp_ratio=self.mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], norm_layer=norm_layer, downsample=PatchMerging if (i_layer < self.num_layers - 1) else None, use_checkpoint=use_checkpoint) self.layers.append(layer)

# 1. Patch Partition : 이미지를 패치로 분할하고, 패치에 포지션 임베딩을 적용여부를 설정 self.patch_embed = PatchEmbed( img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, norm_layer=norm_layer if self.patch_norm else None) num_patches = self.patch_embed.num_patches patches_resolution = self.patch_embed.patches_resolution self.patches_resolution = patches_resolution # Absolute Position Embedding을 원할 경우 (디폴트는 Relative Position bias) if self.ape: self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) trunc_normal_(self.absolute_pos_embed, std=.02) # drop-out 설정 self.pos_drop = nn.Dropout(p=drop_rate) # drop path rate 계산 dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]

class PatchEmbed(nn.Module): def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): super().__init__() img_size = to_2tuple(img_size) # 이미지 크기를 튜플 형태로 변환 patch_size = to_2tuple(patch_size) # 패치 크기를 튜플 형태로 변환 # 각 차원의 패치 개수를 계산 (H/4, W/4) patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] self.img_size = img_size # 이미지 크기 변수 설정 self.patch_size = patch_size # 패치 크기 변수 설정 self.patches_resolution = patches_resolution # 패치 해상도 변수 설정 # 전체 패치 수를 계산하여 변수 설정 self.num_patches = patches_resolution[0] * patches_resolution[1] self.in_chans = in_chans # 입력 채널 수 self.embed_dim = embed_dim # 임베딩 차원 수 # 컨볼루션 레이어를 이용한 패치 Linear Projection Layer (48→ C) self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) if norm_layer is not None: # 정규화 레이어가 있는 경우에 대한 예외처리 self.norm = norm_layer(embed_dim) # 정규화 레이어를 초기화 else: self.norm = None # 정규화 레이어 없음 def forward(self, x): B, C, H, W = x.shape assert H == self.img_size[0] and W == self.img_size[1], \ f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." x = self.proj(x).flatten(2).transpose(1, 2) if self.norm is not None: x = self.norm(x) return x def flops(self): Ho, Wo = self.patches_resolution flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1]) if self.norm is not None: flops += Ho * Wo * self.embed_dim return flops

class BasicLayer(nn.Module): def __init__(self, dim, input_resolution, depth, num_heads, window_size, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False): super().__init__() self.dim = dim # 입력 차원 self.input_resolution = input_resolution # 입력 해상도 self.depth = depth # 레이어 깊이 self.use_checkpoint = use_checkpoint # 체크포인팅 사용 여부 # for문을 사용해서 한 스테이지 내의 depth(블록 수)만큼의 Swin Transformer 블록 구축 self.blocks = nn.ModuleList([ SwinTransformerBlock(dim=dim, input_resolution=input_resolution, num_heads=num_heads, window_size=window_size, shift_size=0 if (i % 2 == 0) else window_size // 2, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop, attn_drop=attn_drop, drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, norm_layer=norm_layer) for i in range(depth)]) # Stage 2, 3, 4의 Patch Merging Layer if downsample is not None: self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) else: self.downsample = None def forward(self, x): for blk in self.blocks: if self.use_checkpoint: x = checkpoint.checkpoint(blk, x) else: x = blk(x) if self.downsample is not None: x = self.downsample(x) return x def extra_repr(self) -> str: return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" def flops(self): flops = 0 for blk in self.blocks: flops += blk.flops() if self.downsample is not None: flops += self.downsample.flops() return flops

class SwinTransformerBlock(nn.Module): def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm): super().__init__() self.dim = dim # 입력 차원 self.input_resolution = input_resolution # 입력 해상도 self.num_heads = num_heads # 어텐션 헤드 수 self.window_size = window_size # 윈도우 크기 self.shift_size = shift_size # 시프트 크기 self.mlp_ratio = mlp_ratio # MLP 차원 비율 if min(self.input_resolution) <= self.window_size: # 입력 해상도보다 윈도우 크기가 크거나 같은 경우 윈도우 파티셔닝을 하지 않음 self.shift_size = 0 self.window_size = min(self.input_resolution) assert 0 <= self.shift_size < self.window_size, "shift_size must be in 0-window_size" # 1. 첫 번째 정규화 레이어 self.norm1 = norm_layer(dim) # 2. WA 레이어 self.attn = WindowAttention( dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) # 윈도우 어텐션 self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() # 드롭 패스 # 3. 두 번째 정규화 레이어 self.norm2 = norm_layer(dim) # 4. MLP 레이어 mlp_hidden_dim = int(dim * mlp_ratio) self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) # SW-MSA를 위한 어텐션 마스크 계산 if self.shift_size > 0: H, W = self.input_resolution img_mask = torch.zeros((1, H, W, 1)) # 이미지 마스크 초기화 h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 mask_windows = window_partition(img_mask, self.window_size) mask_windows = mask_windows.view(-1, self.window_size * self.window_size) attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None self.register_buffer("attn_mask", attn_mask) # 어텐션 마스크 버퍼 등록 def forward(self, x): '''SwinTransformerBlock에서 어텐션과 피드포워드 네트워크를 처리하는 forward 함수''' H, W = self.input_resolution B, L, C = x.shape assert L == H * W, "input feature has wrong size" shortcut = x # 잔차연결을 위한 입력값을 shortcut에 저장 x = self.norm1(x) # 1번째 정규화 레이어 norm1 적용 x = x.view(B, H, W, C) # 순환 시프트 적용 if self.shift_size > 0: # shift_size가 0보다 큰 경우 # 입력 텐서를 지정된 크기만큼 순환 shift (SW-MSA를 위한 과정) shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) else: shifted_x = x # 윈도우 분할 - 순환 시프트된 텐서를 윈도우 단위로 분할 (window_size, window_size, C) x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C # W-MSA/SW-MSA - 분할된 윈도우들에 대해 어텐션 메커니즘 적용 attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C # 윈도우 병합 attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) # window_reverse 함수로 어텐션을 적용한 윈도우들을 원래 이미지 형태로 병합 shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C # 순환 시프트 복원 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) # FFN x = shortcut + self.drop_path(x) x = x + self.drop_path(self.mlp(self.norm2(x))) return x def extra_repr(self) -> str: return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" def flops(self): flops = 0 H, W = self.input_resolution # norm1 flops += self.dim * H * W # W-MSA/SW-MSA nW = H * W / self.window_size / self.window_size flops += nW * self.attn.flops(self.window_size * self.window_size) # mlp flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio # norm2 flops += self.dim * H * W return flops

class WindowAttention(nn.Module): def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): super().__init__() self.dim = dim # 입력 차원 self.window_size = window_size # 윈도우 크기 (Wh, Ww) self.num_heads = num_heads # 어텐션 헤드 수 head_dim = dim // num_heads # 각 헤드의 차원 self.scale = qk_scale or head_dim ** -0.5 # 스케일링 인자 # 1. relative position 테이블 정의 self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 상대 위치에 따른 바이어스 값 # 2. 각 토큰 간의 relative position 인덱스를 계산 coords_h = torch.arange(self.window_size[0]) coords_w = torch.arange(self.window_size[1]) # 좌표 격자 생성 coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 좌표 평탄화 coords_flatten = torch.flatten(coords, 1) # 상대 좌표 계산 relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 차원 재배열 relative_coords = relative_coords.permute(1, 2, 0).contiguous() # index 값이 0부터 시작하도록 조정 relative_coords[:, :, 0] += self.window_size[0] - 1 #x-axis matrix relative_coords[:, :, 1] += self.window_size[1] - 1 #y-axis matrix relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 #x-axis matrix : 2*window_size -1 곱하기 relative_position_index = relative_coords.sum(-1) # 최종 상대 위치 인덱스 self.register_buffer("relative_position_index", relative_position_index) # QKV(쿼리, 키, 값) 선형 변환을 정의 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) # 어텐션 드롭아웃 self.proj = nn.Linear(dim, dim) # 결과 프로젝션 레이어 self.proj_drop = nn.Dropout(proj_drop) # 프로젝션 드롭아웃 # relative position bias table 초기화 trunc_normal_(self.relative_position_bias_table, std=.02) # 소프트맥스 레이어로 최종 어텐션 확률 계산 self.softmax = nn.Softmax(dim=-1) def forward(self, x, mask=None): """ parameter : x: 입력 특징 벡터 (num_windows*B, N, C) mask: 어텐션 마스크 (num_windows, Wh*Ww, Wh*Ww) 또는 None """ B_, N, C = x.shape qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # 쿼리, 키, 값 분리 q = q * self.scale # 쿼리 스케일링 attn = (q @ k.transpose(-2, -1)) # 어텐션 스코어 계산 # relative position bias 추가 relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) relative_position_bias = relative_position

class Mlp(nn.Module): def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features) self.act = act_layer() self.fc2 = nn.Linear(hidden_features, out_features) self.drop = nn.Dropout(drop) def forward(self, x): x = self.fc1(x) # fully-connected layer x = self.act(x) # GeLU 활성화함수 x = self.drop(x) # drop-out x = self.fc2(x) # fully-connected layer x = self.drop(x) # GeLU 활성화함수 return x

# SW-MSA를 위한 어텐션 마스크 계산 if self.shift_size > 0: H, W = self.input_resolution # 이미지 마스크 초기화 img_mask = torch.zeros((1, H, W, 1)) # 너비, 높이에 대한 슬라이스 정의 (window_size, shift_size) h_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) w_slices = (slice(0, -self.window_size), slice(-self.window_size, -self.shift_size), slice(-self.shift_size, None)) # 각 윈도우 위치마다 상대적 위치 인덱스(cnt) 할당 cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, h, w, :] = cnt cnt += 1 # window_partiton 함수로 이미지 마스크를 윈도우 크기에 맞게 분할 mask_windows = window_partition(img_mask, self.window_size) mask_windows = mask_windows.view(-1, self.window_size * self.window_size) # 어텐션 마스크 생성 # 윈도우의 값을 1차원으로 평활화 attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) # 서로 인접한 부분은 0, 인접하지 않은 부분은 -100으로 마스크 생성 attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) else: attn_mask = None

def forward(self, x): '''SwinTransformerBlock에서 어텐션과 피드포워드 네트워크를 처리하는 forward 함수''' H, W = self.input_resolution B, L, C = x.shape assert L == H * W, "input feature has wrong size" shortcut = x # 잔차연결을 위한 입력값을 shortcut에 저장 x = self.norm1(x) # 1번째 정규화 레이어 norm1 적용 x = x.view(B, H, W, C) # 1. 순환 시프트(Cyclie shift) 적용 if self.shift_size > 0: # shift_size가 0보다 큰 경우 # 입력 텐서를 지정된 크기만큼 순환 shift (SW-MSA를 위한 과정) shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) else: shifted_x = x # 윈도우 분할 - 순환 시프트된 텐서를 윈도우 단위로 분할 (window_size, window_size, C) x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C # 2. W-MSA/SW-MSA - 분할된 윈도우들에 대해 어텐션 메커니즘 적용 attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C # 3. 윈도우 병합 attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) # window_reverse 함수로 어텐션을 적용한 윈도우들을 원래 이미지 형태로 병합 shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C # 4. 순환 시프트 복원(reverse cyclic shift) 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) # FFN x = shortcut + self.drop_path(x) # shortcut에 드롭패스 적용결과 추가 # 2번째 normalization 적용 후 mlp 수행 x = x + self.drop_path(self.mlp(self.norm2(x)))