Update metaformers.py

pull/1647/head
Fredo Guan 2 years ago
parent dd57cde1bc
commit 729533e966

@ -441,20 +441,13 @@ class MetaFormerBlock(nn.Module):
x = self.res_scale1(x) + \ x = self.res_scale1(x) + \
self.layer_scale1( self.layer_scale1(
self.drop_path1( self.drop_path1(
self.token_mixer( self.token_mixer(self.norm1(x))
self.norm1(
x.permute(0, 3, 1, 2)
).permute(0, 2, 3, 1)
)
) )
) )
x = self.res_scale2(x) + \ x = self.res_scale2(x) + \
self.layer_scale2( self.layer_scale2(
self.drop_path2( self.drop_path2(
self.mlp(self.norm2( self.mlp(self.norm2(x))
x.permute(0, 3, 1, 2)
)#.permute(0, 2, 3, 1)
)
) )
) )
#x = x.view(B, C, H, W) #x = x.view(B, C, H, W)
@ -915,10 +908,10 @@ def poolformerv1_s24(pretrained=False, **kwargs):
dims=[64, 128, 320, 512], dims=[64, 128, 320, 512],
downsample_norm=None, downsample_norm=None,
token_mixers=Pooling, token_mixers=Pooling,
mlp_fn=partial(nn.Conv2d, kernel_size=1), mlp_fn=partial(Conv2dChannelsLast, kernel_size=1),
mlp_act=nn.GELU, mlp_act=nn.GELU,
mlp_bias=True, mlp_bias=True,
norm_layers=GroupNorm1, norm_layers=partial(LayerNormGeneral, normalized_dim=(1, 2, 3), eps=1e-6, bias=True),
layer_scale_init_values=1e-5, layer_scale_init_values=1e-5,
res_scale_init_values=None, res_scale_init_values=None,
**kwargs) **kwargs)
@ -931,10 +924,10 @@ def poolformerv1_s36(pretrained=False, **kwargs):
dims=[64, 128, 320, 512], dims=[64, 128, 320, 512],
downsample_norm=None, downsample_norm=None,
token_mixers=Pooling, token_mixers=Pooling,
mlp_fn=partial(nn.Conv2d, kernel_size=1), mlp_fn=partial(Conv2dChannelsLast, kernel_size=1),
mlp_act=nn.GELU, mlp_act=nn.GELU,
mlp_bias=True, mlp_bias=True,
norm_layers=GroupNorm1, norm_layers=partial(LayerNormGeneral, normalized_dim=(1, 2, 3), eps=1e-6, bias=True),
layer_scale_init_values=1e-6, layer_scale_init_values=1e-6,
res_scale_init_values=None, res_scale_init_values=None,
**kwargs) **kwargs)
@ -947,10 +940,10 @@ def poolformerv1_m36(pretrained=False, **kwargs):
dims=[96, 192, 384, 768], dims=[96, 192, 384, 768],
downsample_norm=None, downsample_norm=None,
token_mixers=Pooling, token_mixers=Pooling,
mlp_fn=partial(nn.Conv2d, kernel_size=1), mlp_fn=partial(Conv2dChannelsLast, kernel_size=1),
mlp_act=nn.GELU, mlp_act=nn.GELU,
mlp_bias=True, mlp_bias=True,
norm_layers=GroupNorm1, norm_layers=partial(LayerNormGeneral, normalized_dim=(1, 2, 3), eps=1e-6, bias=True),
layer_scale_init_values=1e-6, layer_scale_init_values=1e-6,
res_scale_init_values=None, res_scale_init_values=None,
**kwargs) **kwargs)
@ -963,10 +956,10 @@ def poolformerv1_m48(pretrained=False, **kwargs):
dims=[96, 192, 384, 768], dims=[96, 192, 384, 768],
downsample_norm=None, downsample_norm=None,
token_mixers=Pooling, token_mixers=Pooling,
mlp_fn=partial(nn.Conv2d, kernel_size=1), mlp_fn=partial(Conv2dChannelsLast, kernel_size=1),
mlp_act=nn.GELU, mlp_act=nn.GELU,
mlp_bias=True, mlp_bias=True,
norm_layers=GroupNorm1, norm_layers=partial(LayerNormGeneral, normalized_dim=(1, 2, 3), eps=1e-6, bias=True),
layer_scale_init_values=1e-6, layer_scale_init_values=1e-6,
res_scale_init_values=None, res_scale_init_values=None,
**kwargs) **kwargs)

Loading…
Cancel
Save