mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Update metaformers.py

Browse Source
This commit is contained in:
Fredo Guan 2023-01-08 00:12:52 -08:00
parent 926d886527
commit 7f149f31d4
1 changed files with 14 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
timm/models/metaformers.py
View File

@ -521,14 +521,11 @@ class MetaFormerBlock(nn.Module):
norm_layer=nn.LayerNorm,
drop=0., drop_path=0.,
layer_scale_init_value=None,
res_scale_init_value=None,
downsample = nn.Identity()
res_scale_init_value=None
):
super().__init__()
self.downsample = downsample
self.norm1 = norm_layer(dim)
self.token_mixer = token_mixer(dim=dim, drop=drop)
self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
@ -546,7 +543,6 @@ class MetaFormerBlock(nn.Module):
if res_scale_init_value else nn.Identity()
def forward(self, x):
x = self.downsample(x)
x = self.res_scale1(x) + \
self.layer_scale1(
self.drop_path1(
@ -653,18 +649,19 @@ class MetaFormer(nn.Module):
stages = nn.ModuleList() # each stage consists of multiple metaformer blocks
cur = 0
for i in range(num_stage):
stage = nn.Sequential(*[MetaFormerBlock(
stage = nn.Sequential(OrderedDict[
('downsample', downsample_layers[i]),
('blocks', nn.Sequential(*[MetaFormerBlock(
dim=dims[i],
token_mixer=token_mixers[i],
mlp=mlps[i],
norm_layer=norm_layers[i],
drop_path=dp_rates[cur + j],
layer_scale_init_value=layer_scale_init_values[i],
res_scale_init_value=res_scale_init_values[i],
downsample = downsample_layers[i]
) for j in range(depths[i])]
res_scale_init_value=res_scale_init_values[i]
) for j in range(depths[i])])
)]
)
stages.append(stage)
cur += depths[i]