cngthnh/Thesis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

try modify swin

Browse Source
This commit is contained in:
thanhvc3 2024-04-29 17:06:55 +07:00
parent 48669c72f4
commit 5494206a04
1 changed files with 39 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
models.py
View File

@ -547,6 +547,7 @@ class FouriER(torch.nn.Module):
down_patch_size=3 down_patch_size=3
down_stride=2 down_stride=2
down_pad=1 down_pad=1
window_size = 4
num_classes=self.p.embed_dim num_classes=self.p.embed_dim
for i in range(len(layers)): for i in range(len(layers)):
stage = basic_blocks(embed_dims[i], i, layers, stage = basic_blocks(embed_dims[i], i, layers,
@ -556,7 +557,8 @@ class FouriER(torch.nn.Module):
drop_path_rate=drop_path_rate, drop_path_rate=drop_path_rate,
use_layer_scale=use_layer_scale, use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value, layer_scale_init_value=layer_scale_init_value,
num_heads=num_heads[i]) num_heads=num_heads[i], input_resolution=(image_h // (2**i), image_w // (2**i)),
window_size=window_size, shift_size=0 if (i % 2 == 0) else window_size // 2)
network.append(stage) network.append(stage)
if i >= len(layers) - 1: if i >= len(layers) - 1:
break break
@ -734,7 +736,7 @@ def basic_blocks(dim, index, layers,
pool_size=3, mlp_ratio=4., pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm, act_layer=nn.GELU, norm_layer=GroupNorm,
drop_rate=.0, drop_path_rate=0., drop_rate=.0, drop_path_rate=0.,
use_layer_scale=True, layer_scale_init_value=1e-5, num_heads = 4): use_layer_scale=True, layer_scale_init_value=1e-5, num_heads = 4, input_resolution = None, window_size = 4, shift_size = 2):
""" """
generate PoolFormer blocks for a stage generate PoolFormer blocks for a stage
return: PoolFormer blocks return: PoolFormer blocks
@ -749,7 +751,8 @@ def basic_blocks(dim, index, layers,
drop=drop_rate, drop_path=block_dpr, drop=drop_rate, drop_path=block_dpr,
use_layer_scale=use_layer_scale, use_layer_scale=use_layer_scale,
layer_scale_init_value=layer_scale_init_value, layer_scale_init_value=layer_scale_init_value,
num_heads=num_heads num_heads=num_heads, input_resolution = input_resolution,
window_size=window_size, shift_size=shift_size
)) ))
blocks = nn.Sequential(*blocks) blocks = nn.Sequential(*blocks)
@ -933,14 +936,16 @@ class PoolFormerBlock(nn.Module):
def __init__(self, dim, pool_size=3, mlp_ratio=4., def __init__(self, dim, pool_size=3, mlp_ratio=4.,
act_layer=nn.GELU, norm_layer=GroupNorm, act_layer=nn.GELU, norm_layer=GroupNorm,
drop=0., drop_path=0., num_heads=4, drop=0., drop_path=0., num_heads=4,
use_layer_scale=True, layer_scale_init_value=1e-5): use_layer_scale=True, layer_scale_init_value=1e-5, input_resolution = None, window_size = 4, shift_size = 2):
super().__init__() super().__init__()
self.norm1 = norm_layer(dim) self.norm1 = norm_layer(dim)
#self.token_mixer = Pooling(pool_size=pool_size) #self.token_mixer = Pooling(pool_size=pool_size)
# self.token_mixer = FNetBlock() # self.token_mixer = FNetBlock()
self.window_size = 4 self.window_size = window_size
self.shift_size = shift_size
self.input_resolution = input_resolution
self.attn_mask = None self.attn_mask = None
self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, attn_drop=0.1, proj_drop=0.2) self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, attn_drop=0.1, proj_drop=0.2)
self.norm2 = norm_layer(dim) self.norm2 = norm_layer(dim)
@ -958,6 +963,31 @@ class PoolFormerBlock(nn.Module):
self.layer_scale_2 = nn.Parameter( self.layer_scale_2 = nn.Parameter(
layer_scale_init_value * torch.ones((dim)), requires_grad=True) layer_scale_init_value * torch.ones((dim)), requires_grad=True)
if self.shift_size > 0:
# calculate attention mask for SW-MSA
H, W = self.input_resolution
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))
else:
attn_mask = None
self.register_buffer("attn_mask", attn_mask)
def forward(self, x): def forward(self, x):
B, C, H, W = x.shape B, C, H, W = x.shape
x_windows = window_partition(x, self.window_size) x_windows = window_partition(x, self.window_size)
@ -965,6 +995,10 @@ class PoolFormerBlock(nn.Module):
attn_windows = self.token_mixer(x_windows, mask=self.attn_mask) attn_windows = self.token_mixer(x_windows, mask=self.attn_mask)
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
x_attn = window_reverse(attn_windows, self.window_size, H, W) x_attn = window_reverse(attn_windows, self.window_size, H, W)
if self.shift_size > 0:
x = torch.roll(x_attn, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
x = x_attn
if self.use_layer_scale: if self.use_layer_scale:
x = x + self.drop_path( x = x + self.drop_path(
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)