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

Fix drop/drop_path arg on MLP-Mixer model. Fix #641

Browse Source
This commit is contained in:
Ross Wightman 2021-05-16 08:31:52 -07:00
parent 7077f16c6a
commit e7f0db8664
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
timm/models/mlp_mixer.py
View File

@ -96,8 +96,8 @@ class MlpMixer(nn.Module):
mlp_layer=Mlp, mlp_layer=Mlp,
norm_layer=partial(nn.LayerNorm, eps=1e-6), norm_layer=partial(nn.LayerNorm, eps=1e-6),
act_layer=nn.GELU, act_layer=nn.GELU,
drop=0., drop_rate=0.,
drop_path=0., drop_path_rate=0.,
nlhb=False, nlhb=False,
): ):
super().__init__() super().__init__()
@ -108,7 +108,7 @@ class MlpMixer(nn.Module):
self.blocks = nn.Sequential(*[ self.blocks = nn.Sequential(*[
MixerBlock( MixerBlock(
hidden_dim, self.stem.num_patches, tokens_dim, channels_dim, hidden_dim, self.stem.num_patches, tokens_dim, channels_dim,
mlp_layer=mlp_layer, norm_layer=norm_layer, act_layer=act_layer, drop=drop, drop_path=drop_path) mlp_layer=mlp_layer, norm_layer=norm_layer, act_layer=act_layer, drop=drop_rate, drop_path=drop_path_rate)
for _ in range(num_blocks)]) for _ in range(num_blocks)])
self.norm = norm_layer(hidden_dim) self.norm = norm_layer(hidden_dim)
self.head = nn.Linear(hidden_dim, self.num_classes) # zero init self.head = nn.Linear(hidden_dim, self.num_classes) # zero init