- Key1 | LAM으로 Transformer 아키텍처의 문제점을 논리적으로 분석
- Key2 | 더 많은 이미지 픽셀을 활성화하기 위해 channel-attention 사용
- Key3 | Transformer (self-attention schema) 사용
- Key4 | window partition mechanism의 문제점을 해소하기 위해 overlapping cross attention module 도입
1. Introduction
Transformer는 이미지 초해상화 분야에서도 좋은 성과를 보이고 있습니다. 그러나 초해상화에 비교적 작은 영역만을 활용하기 때문에 성능 향상의 여지가 있습니다. 다시 말해 고화질 이미지를 생성할 때 더 많은 픽셀을 활성화한다면 더 좋은 성능을 보일 수도 있다는 것입니다. 이를 위해 본 논문에서는 HAT : Hybrid Attention Transformer 아키텍처를 소개합니다. HAT에는 channel-attention과 self-attention이 사용된 것이 특징입니다.
Transformer 구조가 CNN보다 뛰어난 이유는 여전히 확실하지 않습니다. 직관적으로는 Trnasformer의 self-attention 모듈이 long-range의 픽셀 정보를 활용하기 때문입니다. Fig. 2를 보면 CNN 기반 모델인 EDSR과 RCAN 사이에서는 활성화 픽셀의 범위가 넓은 RCAN이 더 좋은 초해상화 성능을 보입니다. 하지만 Transformer를 차용한 SwinIR의 초해상화 결과를 분석하면 CNN 기반의 RCAN 아키텍처보다 픽셀 활성도가 뒤쳐졌습니다. 그럼에도 초해상화의 정량적 성능은 SwinIR이 더 좋은데 이는 Transformer의 기본적인 성능(논문에서는 stronger mapping capability라고 표현)이 더 뛰어나기 때문이라고 할 수 있습니다. 하지만 PSNR/SSIM 점수가 무색하게 SwinIR는 파란색 박스 영역을 제대로 복원하지 못했습니다. 논문 저자들은 RCAN과 SwinIR의 차이가 channel attention 모듈에서 기인한다고 봤습니다.
Fig. 2의 (b)에는 SwinIR의 중간층 특징맵에는 격자 무늬의 아티팩트가 있는 것을 발견했습니다. 원인은 window sliding method이고 이를 해결하기 위해 overlapping cross-attention 모듈을 도입했습니다.
2. Network Architecture
HAT의 아키텍처는 Shallow Feature Extraction, Deep Feature Extraction, Image Reconstruction으로 구성됩니다. Shallow Feature Extraction은 한 개의 합성곱층으로 픽셀을 토큰화하는 역할을 합니다. RGB 이미지를 Shallow Feature Extraction에 입력해 C 차원의 임베딩 벡터(특징맵)를 얻을 수 있습니다. 초기에 합성곱층을 활용하면 모델이 더 좋은 이미지 표현을 배우고, 안정적으로 최적화하는데 도움이 됩니다.
Deep Feature Extraction은 N개의 RHAG(Residual Hybrid Attention Groups)와 3x3 합성곱층 한 개로 구성됩니다. 마지막에 사용하는 합성곱층은 deep features의 정보를 취합하는 데 도움을 줍니다. 그리고 나서는 Shallow feature와 Deep feature를 Elementwise-sum 합니다.
Reconstruction Module에서는 pixel-shuffle method를 이용해 feature를 Upsample 합니다. 손실함수로는 L1 loss를 사용합니다.
Residual Hybrid Attention Groups (RHAG)
- Fi-1, 0 : i번째 RHAG모듈의 입력 특징맵
- Fi-1, j : i번째 RHAG모듈의 j번째 HAB 모듈의 결과 특징맵
- Fi : i번째 RHAG모듈의 결과 특징맵
RHAG에서는 먼저 입력 특징맵을 HAB(Hybrid Attention Block) 모듈 여러 개와 OCAB(Overlapping Cross-Attention Block) 모듈, 합성곱 층을 통과시킨 특징맵과 입력 특징맵을 elementwise-sum해 결과를 냅니다.
RHAG 코드
더보기
class RHAG(nn.Module):
"""Residual Hybrid Attention Group (RHAG).
Args:
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
depth (int): Number of blocks.
num_heads (int): Number of attention heads.
window_size (int): Local window size.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
img_size: Input image size.
patch_size: Patch size.
resi_connection: The convolutional block before residual connection.
"""
def __init__(self,
dim,
input_resolution,
depth,
num_heads,
window_size,
compress_ratio,
squeeze_factor,
conv_scale,
overlap_ratio,
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,
img_size=224,
patch_size=4,
resi_connection='1conv'):
super(RHAG, self).__init__()
self.dim = dim
self.input_resolution = input_resolution
self.residual_group = AttenBlocks(
dim=dim,
input_resolution=input_resolution,
depth=depth,
num_heads=num_heads,
window_size=window_size,
compress_ratio=compress_ratio,
squeeze_factor=squeeze_factor,
conv_scale=conv_scale,
overlap_ratio=overlap_ratio,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop=drop,
attn_drop=attn_drop,
drop_path=drop_path,
norm_layer=norm_layer,
downsample=downsample,
use_checkpoint=use_checkpoint)
if resi_connection == '1conv':
self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
elif resi_connection == 'identity':
self.conv = nn.Identity()
self.patch_embed = PatchEmbed(
img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)
self.patch_unembed = PatchUnEmbed(
img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)
def forward(self, x, x_size, params):
return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size, params), x_size))) + x출처 : https://github.com/XPixelGroup/HAT/blob/main/hat/archs/hat_arch.py
AttenBlocks 코드
더보기
class AttenBlocks(nn.Module):
""" A series of attention blocks for one RHAG.
Args:
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
depth (int): Number of blocks.
num_heads (int): Number of attention heads.
window_size (int): Local window size.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
"""
def __init__(self,
dim,
input_resolution,
depth,
num_heads,
window_size,
compress_ratio,
squeeze_factor,
conv_scale,
overlap_ratio,
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
# build blocks
self.blocks = nn.ModuleList([
HAB(
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,
compress_ratio=compress_ratio,
squeeze_factor=squeeze_factor,
conv_scale=conv_scale,
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)
])
# OCAB
self.overlap_attn = OCAB(
dim=dim,
input_resolution=input_resolution,
window_size=window_size,
overlap_ratio=overlap_ratio,
num_heads=num_heads,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
mlp_ratio=mlp_ratio,
norm_layer=norm_layer
)
# 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, x_size, params):
for blk in self.blocks:
x = blk(x, x_size, params['rpi_sa'], params['attn_mask'])
x = self.overlap_attn(x, x_size, params['rpi_oca'])
if self.downsample is not None:
x = self.downsample(x)
return x
PatchEmbed, PatchUnEmbed 코드
더보기
class PatchEmbed(nn.Module):
r""" Image to Patch Embedding
Args:
img_size (int): Image size. Default: 224.
patch_size (int): Patch token size. Default: 4.
in_chans (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
norm_layer (nn.Module, optional): Normalization layer. Default: None
"""
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) # (224, 224)
patch_size = to_2tuple(patch_size) # (4, 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
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
self.norm = None
def forward(self, x):
x = x.flatten(2).transpose(1, 2) # b c, Ph, Pw -> b Ph*Pw c
if self.norm is not None:
x = self.norm(x)
return x
class PatchUnEmbed(nn.Module):
r""" Image to Patch Unembedding
Args:
img_size (int): Image size. Default: 224.
patch_size (int): Patch token size. Default: 4.
in_chans (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
norm_layer (nn.Module, optional): Normalization layer. Default: None
"""
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)
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
def forward(self, x, x_size):
x = x.transpose(1, 2).contiguous().view(x.shape[0], self.embed_dim, x_size[0], x_size[1]) # b embed_dim Ph Pw
return xHybrid Attention Block (HAB)
- LN : Layer Normalization
HAB 모듈은 크게 Channel Attention Block(CAB)와 Window-based Multi-head Self-Attention(W-MSA)로 구성된다. Shifted W-MSA(SW-MSA)를 번갈아 사용한다. 최적화 과정에서 CAB와 MSA가 충돌하는 것을 방지하기 위해 CAB에는 작은 상수(α)를 곱해준다. 개별 픽셀은 토큰으로 임베딩한다. 즉 patch size가 1이다.
HAB 코드
더보기
def window_partition(x, window_size):
"""
Args:
x: (b, h, w, c)
window_size (int): window size
Returns:
windows: (num_windows*b, window_size, window_size, c)
"""
b, h, w, c = x.shape
x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)
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 (int): Window size
h (int): Height of image
w (int): Width of image
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
class HAB(nn.Module):
r""" Hybrid Attention Block.
Args:
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
num_heads (int): Number of attention heads.
window_size (int): Window size.
shift_size (int): Shift size for SW-MSA.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
"""
def __init__(self,
dim,
input_resolution,
num_heads,
window_size=7,
shift_size=0,
compress_ratio=3,
squeeze_factor=30,
conv_scale=0.01, #alpha
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
if min(self.input_resolution) <= self.window_size:
# if window size is larger than input resolution, we don't partition windows
self.shift_size = 0
self.window_size = min(self.input_resolution)
assert 0 <= self.shift_size < self.window_size, 'shift_size must in 0-window_size'
self.norm1 = norm_layer(dim)
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.conv_scale = conv_scale
self.conv_block = CAB(num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor)
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x, x_size, rpi_sa, attn_mask):
h, w = x_size
b, _, c = x.shape
# assert seq_len == h * w, "input feature has wrong size"
shortcut = x
x = self.norm1(x)
x = x.view(b, h, w, c)
# Conv_X
conv_x = self.conv_block(x.permute(0, 3, 1, 2))
conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c)
# cyclic shift
if self.shift_size > 0:
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
attn_mask = attn_mask
else:
shifted_x = x
attn_mask = None
# partition windows
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 (to be compatible for testing on images whose shapes are the multiple of window size
attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask) # window 단위로 self-attention 진행
# merge windows
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)
shifted_x = window_reverse(attn_windows, self.window_size, h, w) # b h' w' c
# reverse cyclic shift
if self.shift_size > 0:
attn_x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
attn_x = shifted_x
attn_x = attn_x.view(b, h * w, c)
# FFN
x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
DropPath 코드
더보기
def drop_path(x, drop_prob: float = 0., training: bool = False):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
"""
if drop_prob == 0. or not training:
return x
keep_prob = 1 - drop_prob
shape = (x.shape[0], ) + (1, ) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
random_tensor.floor_() # binarize
output = x.div(keep_prob) * random_tensor
return output
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
"""
def __init__(self, drop_prob=None):
super(DropPath, self).__init__()
self.drop_prob = drop_prob
def forward(self, x):
return drop_path(x, self.drop_prob, self.training)
Mlp 코드
더보기
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)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
(S)W-MSA
- d : dimension of query / key
- B : relative position encoding
self-attention을 계산하기 전에 H x W x C을 HW/M²개 윈도우로 나눕니다. 이때 윈도우의 크기는 M x M이 됩니다. 그리고 나서 개별 윈도우 내에서 self-attention을 진행합니다.
또한 이웃하는 윈도우의 연결성을 위해 윈도우 사이즈의 반만큼 이동시키는 전략을 사용합니다. (HAB 코드에서 설정)
W-MSA 코드
더보기
class WindowAttention(nn.Module):
r""" Window based multi-head self attention (W-MSA) module with relative position bias.
It supports both of shifted and non-shifted window.
Args:
dim (int): Number of input channels.
window_size (tuple[int]): The height and width of the window.
num_heads (int): Number of attention heads.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
"""
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
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
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)
trunc_normal_(self.relative_position_bias_table, std=.02) # function reference : https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/archs/arch_util.py
self.softmax = nn.Softmax(dim=-1)
def forward(self, x, rpi, mask=None):
"""
Args:
x: input features with shape of (num_windows*b, n, c)
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or 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] # make torchscript happy (cannot use tensor as tuple)
q = q * self.scale
attn = (q @ k.transpose(-2, -1))
relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
attn = attn + relative_position_bias.unsqueeze(0)
if mask is not None:
nw = mask.shape[0]
attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, n, n)
attn = self.softmax(attn)
else:
attn = self.softmax(attn)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(b_, n, c)
x = self.proj(x)
x = self.proj_drop(x)
return x출처 : https://github.com/XPixelGroup/HAT/blob/main/hat/archs/hat_arch.py
CAB
CAB는 두 개의 합성곱층과 channel-attention module로 구성됩니다. Transformer 구조의 임베딩 벡터(토큰)의 채널(C)은 매우 크기 때문에 그대로 사용하면 계산 비용이 너무 많이 듭니다. 따라서 첫 합성곱층에서 상수 β로 나눈 것 만큼 채널을 압축합니다. 첫 합성곱층의 결과 특징맵의 채널 수는 C/β가 됩니다. 합성곱층 사이에는 활성화 함수로 GELU를 사용합니다. 두 합성곱층을 통과한 feature는 RCAN에서 사용한 channel-attention 모듈에 입력해 결과 특징맵을 얻을 수 있습니다.
CAB 코드
더보기
class ChannelAttention(nn.Module):
"""Channel attention used in RCAN.
Args:
num_feat (int): Channel number of intermediate features.
squeeze_factor (int): Channel squeeze factor. Default: 16.
"""
def __init__(self, num_feat, squeeze_factor=16):
super(ChannelAttention, self).__init__()
self.attention = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),
nn.ReLU(inplace=True),
nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),
nn.Sigmoid())
def forward(self, x):
y = self.attention(x)
return x * y
class CAB(nn.Module):
def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):
super(CAB, self).__init__()
self.cab = nn.Sequential(
nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),
nn.GELU(),
nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),
ChannelAttention(num_feat, squeeze_factor)
)
def forward(self, x):
return self.cab(x)
Overlapping Cross-Attention Block (OCAB)
OCAB 구조는 직접적으로 윈도우 사이의 연결을 강화하기 위한 모듈입니다. 구조는 Swin Transformer와 유사합니다. 차이점은 OCAB에서는 overlapping window partition을 기반으로 self-attention을 한다는 것입니다. overlapping window partition은 query(M)보다 더 큰 사이즈(M0)의 cross window feature로부터 kye/value를 만듭니다.
"Aggregating global features into local vision transformer. (2022)" 논문의 Multi-resolution Overlapped Attention(MOA) 모듈도 overlapping window partition과 비슷하게 동작하지만, MOA는 global attention을 계산하고 OCA는 윈도우 안에서 cross-attention을 계산한다는 차이가 있습니다.
OCAB 코드
더보기
class OCAB(nn.Module):
# overlapping cross-attention block
def __init__(self, dim,
input_resolution,
window_size,
overlap_ratio, # gamma
num_heads,
qkv_bias=True,
qk_scale=None,
mlp_ratio=2,
norm_layer=nn.LayerNorm
):
super().__init__()
self.dim = dim
self.input_resolution = input_resolution
self.window_size = window_size
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim**-0.5
self.overlap_win_size = int(window_size * overlap_ratio) + window_size
self.norm1 = norm_layer(dim)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size),
stride=window_size,
padding=(self.overlap_win_size-window_size)//2)
# define a parameter table of relative position bias
self.relative_position_bias_table = nn.Parameter(
torch.zeros((window_size + self.overlap_win_size - 1) * (window_size + self.overlap_win_size - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
trunc_normal_(self.relative_position_bias_table, std=.02)
self.softmax = nn.Softmax(dim=-1)
self.proj = nn.Linear(dim,dim)
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)
def forward(self, x, x_size, rpi):
h, w = x_size
b, _, c = x.shape
shortcut = x
x = self.norm1(x)
x = x.view(b, h, w, c)
qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2) # 3, b, c, h, w
q = qkv[0].permute(0, 2, 3, 1) # b, h, w, c
kv = torch.cat((qkv[1], qkv[2]), dim=1) # b, 2*c, h, w
# partition windows
q_windows = window_partition(q, self.window_size) # nw*b, window_size, window_size, c
q_windows = q_windows.view(-1, self.window_size * self.window_size, c) # nw*b, window_size*window_size, c
kv_windows = self.unfold(kv) # b, c*w*w, nw
kv_windows = rearrange(kv_windows, 'b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch',
nc=2, ch=c, owh=self.overlap_win_size, oww=self.overlap_win_size).contiguous() # 2, nw*b, ow*ow, c
k_windows, v_windows = kv_windows[0], kv_windows[1] # nw*b, ow*ow, c
b_, nq, _ = q_windows.shape # nq = window_size^2
_, n, _ = k_windows.shape # n = overlap_win_size^2
d = self.dim // self.num_heads
q = q_windows.reshape(b_, nq, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, nq, d
k = k_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d
v = v_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d
q = q * self.scale
attn = (q @ k.transpose(-2, -1)) # (nw*b, nH, nq, d) @ (nw*d, nH, d, n) -> (nw*d, nH, nq, n)
relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(
self.window_size * self.window_size, self.overlap_win_size * self.overlap_win_size, -1) # ws*ws, wse*wse, nH
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, ws*ws, wse*wse
attn = attn + relative_position_bias.unsqueeze(0) # (nw*d, nH, nq, n) + (1, nH, ws*ws, wse*wse)
attn = self.softmax(attn)
attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)
# (nw*d, nH, nq, n) @ (nw*b, nH, n, d) -> (nw*b, nH, nq, d) -> (nw*b=b_, nq , nH*d=self.dim)
# merge windows
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim)
x = window_reverse(attn_windows, self.window_size, h, w) # b h w c
x = x.view(b, h * w, self.dim)
x = self.proj(x) + shortcut
x = x + self.mlp(self.norm2(x))
return x
HAT 코드
더보기
class HAT(nn.Module):
r""" Hybrid Attention Transformer
A PyTorch implementation of : `Activating More Pixels in Image Super-Resolution Transformer`.
Some codes are based on SwinIR.
Args:
img_size (int | tuple(int)): Input image size. Default 64
patch_size (int | tuple(int)): Patch size. Default: 1
in_chans (int): Number of input image channels. Default: 3
embed_dim (int): Patch embedding dimension. Default: 96
depths (tuple(int)): Depth of each Swin Transformer layer.
num_heads (tuple(int)): Number of attention heads in different layers.
window_size (int): Window size. Default: 7
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
drop_rate (float): Dropout rate. Default: 0
attn_drop_rate (float): Attention dropout rate. Default: 0
drop_path_rate (float): Stochastic depth rate. Default: 0.1
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
patch_norm (bool): If True, add normalization after patch embedding. Default: True
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
img_range: Image range. 1. or 255.
upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
"""
def __init__(self,
img_size=64,
patch_size=1,
in_chans=3,
embed_dim=96,
depths=(6, 6, 6, 6),
num_heads=(6, 6, 6, 6),
window_size=7,
compress_ratio=3,
squeeze_factor=30,
conv_scale=0.01,
overlap_ratio=0.5,
mlp_ratio=4.,
qkv_bias=True,
qk_scale=None,
drop_rate=0.,
attn_drop_rate=0.,
drop_path_rate=0.1,
norm_layer=nn.LayerNorm,
ape=False,
patch_norm=True,
use_checkpoint=False,
upscale=2,
img_range=1.,
upsampler='',
resi_connection='1conv',
**kwargs):
super(HAT, self).__init__()
self.window_size = window_size
self.shift_size = window_size // 2
self.overlap_ratio = overlap_ratio
num_in_ch = in_chans
num_out_ch = in_chans
num_feat = 64
self.img_range = img_range
if in_chans == 3:
rgb_mean = (0.4488, 0.4371, 0.4040)
self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)
else:
self.mean = torch.zeros(1, 1, 1, 1)
self.upscale = upscale
self.upsampler = upsampler
# relative position index
relative_position_index_SA = self.calculate_rpi_sa()
relative_position_index_OCA = self.calculate_rpi_oca()
self.register_buffer('relative_position_index_SA', relative_position_index_SA)
self.register_buffer('relative_position_index_OCA', relative_position_index_OCA)
# ------------------------- 1, shallow feature extraction ------------------------- #
self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
# ------------------------- 2, deep feature extraction ------------------------- #
self.num_layers = len(depths)
self.embed_dim = embed_dim
self.ape = ape
self.patch_norm = patch_norm
self.num_features = embed_dim
self.mlp_ratio = mlp_ratio
# split image into non-overlapping patches
self.patch_embed = PatchEmbed(
img_size=img_size,
patch_size=patch_size,
in_chans=embed_dim,
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
# merge non-overlapping patches into image
self.patch_unembed = PatchUnEmbed(
img_size=img_size,
patch_size=patch_size,
in_chans=embed_dim,
embed_dim=embed_dim,
norm_layer=norm_layer if self.patch_norm else None)
# absolute position embedding
if self.ape:
self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
trunc_normal_(self.absolute_pos_embed, std=.02)
self.pos_drop = nn.Dropout(p=drop_rate)
# stochastic depth
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule
# build Residual Hybrid Attention Groups (RHAG)
self.layers = nn.ModuleList()
for i_layer in range(self.num_layers):
layer = RHAG(
dim=embed_dim,
input_resolution=(patches_resolution[0], patches_resolution[1]),
depth=depths[i_layer],
num_heads=num_heads[i_layer],
window_size=window_size,
compress_ratio=compress_ratio,
squeeze_factor=squeeze_factor,
conv_scale=conv_scale,
overlap_ratio=overlap_ratio,
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])], # no impact on SR results
norm_layer=norm_layer,
downsample=None,
use_checkpoint=use_checkpoint,
img_size=img_size,
patch_size=patch_size,
resi_connection=resi_connection)
self.layers.append(layer)
self.norm = norm_layer(self.num_features)
# build the last conv layer in deep feature extraction
if resi_connection == '1conv':
self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
elif resi_connection == 'identity':
self.conv_after_body = nn.Identity()
# ------------------------- 3, high quality image reconstruction ------------------------- #
if self.upsampler == 'pixelshuffle':
# for classical SR
self.conv_before_upsample = nn.Sequential(
nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True))
self.upsample = Upsample(upscale, num_feat)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
self.apply(self._init_weights)
def _init_weights(self, m):
if isinstance(m, nn.Linear):
trunc_normal_(m.weight, std=.02)
if isinstance(m, nn.Linear) and m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
def calculate_rpi_sa(self):
# calculate relative position index for SA
coords_h = torch.arange(self.window_size)
coords_w = torch.arange(self.window_size)
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
relative_coords[:, :, 0] += self.window_size - 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size - 1
relative_coords[:, :, 0] *= 2 * self.window_size - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
return relative_position_index
def calculate_rpi_oca(self):
# calculate relative position index for OCA
window_size_ori = self.window_size
window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size)
coords_h = torch.arange(window_size_ori)
coords_w = torch.arange(window_size_ori)
coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, ws, ws
coords_ori_flatten = torch.flatten(coords_ori, 1) # 2, ws*ws
coords_h = torch.arange(window_size_ext)
coords_w = torch.arange(window_size_ext)
coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, wse, wse
coords_ext_flatten = torch.flatten(coords_ext, 1) # 2, wse*wse
relative_coords = coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None] # 2, ws*ws, wse*wse
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # ws*ws, wse*wse, 2
relative_coords[:, :, 0] += window_size_ori - window_size_ext + 1 # shift to start from 0
relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1
relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1
relative_position_index = relative_coords.sum(-1)
return relative_position_index
def calculate_mask(self, x_size):
# calculate attention mask for SW-MSA
h, w = x_size
img_mask = torch.zeros((1, h, w, 1)) # 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) # nw, window_size, window_size, 1
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))
return attn_mask
@torch.jit.ignore
def no_weight_decay(self):
return {'absolute_pos_embed'}
@torch.jit.ignore
def no_weight_decay_keywords(self):
return {'relative_position_bias_table'}
def forward_features(self, x):
x_size = (x.shape[2], x.shape[3])
# Calculate attention mask and relative position index in advance to speed up inference.
# The original code is very time-cosuming for large window size.
attn_mask = self.calculate_mask(x_size).to(x.device)
params = {'attn_mask': attn_mask, 'rpi_sa': self.relative_position_index_SA, 'rpi_oca': self.relative_position_index_OCA}
x = self.patch_embed(x)
if self.ape:
x = x + self.absolute_pos_embed
x = self.pos_drop(x)
for layer in self.layers:
x = layer(x, x_size, params)
x = self.norm(x) # b seq_len c
x = self.patch_unembed(x, x_size)
return x
def forward(self, x):
self.mean = self.mean.type_as(x)
x = (x - self.mean) * self.img_range
if self.upsampler == 'pixelshuffle':
# for classical SR
x = self.conv_first(x)
x = self.conv_after_body(self.forward_features(x)) + x
x = self.conv_before_upsample(x)
x = self.conv_last(self.upsample(x))
x = x / self.img_range + self.mean
return x
Experiment
same-task pretraining
최근 multi-task pretraining이 각광받고 있는 것 같습니다. multi-task pretraining 모델을 transfer했더니, single-task pretraining 모델을 사용한 것보다 성능이 좋은 경우가 더러 있었단 것이죠. "On efficient transformer and image pre-training for low-level vision." (EDT)논문에서는 transfer할 task와 관련된 여러 task(x2, x3, x4 SR)를 사전학습하는 것이 x4 SR을 위해 x2 SR을 사전학습하거나 x4와 관련 없는 여러 task를 사전학습한 경우보다 효과적이라는 것을 보였습니다. 하지만 논문에서 입증되었다 해도 어떠한 상황에도 적용할 수 있는 진리는 아닌 것 같습니다. 본 논문(HAT)에서는 아래 Table7에 보이는 것과 같이 ImageNet 데이터셋으로 Super Resolution Task만을 사전학습한 다음 target 데이터셋에 transfer한 경우 성능이 더 좋았다고 합니다. 논문 저자는 다음과 같이 쓰며 사전 학습이 효과적인 이유는 사전 학습 과제와 transfer 과제의 연관성 때문일 수 있다고 합니다.
We tend to believe that the reason "why pre-training works" is attributed to the diversity of data instead of the correlation between tasks"
Window Size
SR 과제에서는 입력 픽셀을 더 많이 활성화할 수록 좋은 결과를 얻을 수 있습니다. 이를 위해 직관적으로 window 사이즈를 키우는 방법이 있습니다. "On efficient transformer and image pre-training for low-level vision." (2021) 에서 12x12 크기까지 윈도우 크기에 따른 SR 효과를 실험했습니다. 본 논문에서는 16x16 사이즈까지 실험해 가장 좋은 성능을 얻었습니다. 따라서 HAT 네트워크의 window size도 16으로 설정했습니다.
OCAB / CAB Ablation Study
Overlapping Size
CAB design
Comparison with SOTA Methods
HAT-L은 HAT의 깊이를 두배 늘린 버전으로 RHAG를 6개에서 12개로 늘렸습니다.
Quantitative result
Visual Comparison
Analysis of Model Complexity
- Multiply-Add 연산의 수는 64x64 이미지를 대상으로 카운트
- 사전학습하지 않은 basic 모델 사용
Swin Transormer block에서 window size를 늘리는 경우 파라미터 수, 연산 수가 증가하지만 성능도 개선됩니다.
OCAB는 파라미터와 연산 수를 크게 늘리지 않고 PSNR을 1.1dB 개선했습니다. CAB도 비슷한 성능 개선을 보여줬지만 파라미터 수와 연산 수가 너무 많아지는 문제가 있습니다.
CAB에서 채널 수를 조절하는 β의 영향을 살펴봤습니다. β를 줄일수록 좋은 성능을 보였지만, 계산 비용과 모델 크기를 고려해 β의 기본 설정으로 3으로 정했습니다.
Base Reference인 SwinIR과 HAT을 비교하기 위해 HAT과 비슷한 파라미터 수와 연산 수를 갖도록 SwinIR의 깊이와 너비를 늘렸습니다. 하지만 개인적으로 SwinIR이 이미 최적의 모델 구조를 가졌고, 임의로 구조를 변경해 최적의 구조인 HAT과 비교하는 것은 공정한 비교가 아니라고 생각했습니다.
Training Details
- 저화질 이미지는 고화질 이미지를 bicubic down-sampling해서 얻었습니다.
- 패치 사이즈 : 64 x 64
- data augmentation : random rotation, horizontal flip
- 배치 사이즈 : 32
- Total Iteration : 500ㅏ
- learning rate : 초기에는 2e-4로 설정하고 250K, 400K, 450K, 475K iteration마다 반으로 줄였습니다.
- optimizer : Adam(Beta=(0.9, 0.99)
- x4 SR : x2 SR로 사전학습된 모델을 초기 파라미터로 사용했습니다. 그리고 total iteration과 학습률을 감소시키는 iteration 기준을 절반 수준으로 조정했습니다.
- 사전 학습 : 이미지넷 데이터셋을 전부사용해서 800K iteration 학습했습니다. 학습률은 2e-4로 시작해서 300K, 500K, 650K, 700K, 750K iteration마다 반으로 줄였습니다.
- finetune : DF2K 데이터셋으로 학습했습니다. 초기 학습률은 1e-5로 설정하고 125K, 200K, 230K, 240K iteration마다 반으로 줄였습니다.
More LAM Results
Comparison of Blocking Artifacts
알아볼 것
- LAM
- pixel-shuffle method
- Swin Transformer / SwinIR
'AI > 비전' 카테고리의 다른 글
| [SEGMENTATION] DeepLab v3 | 2017 (0) | 2021.08.17 |
|---|---|
| [SEGMENTATION] DeepLab v2 | 2016 (0) | 2021.08.03 |
| [SEGMENTATION] Attention U-Net: Learning Where to Look for the Pancreas (0) | 2021.07.20 |
| [SEGMENTATION] U-Net++ : Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation | 2018 (0) | 2021.06.28 |
| [SEGMENTATION] DeepLab v1 | 2015 (0) | 2021.06.21 |
