address #16

Author: lucidrains

mlp_mixer_pytorch/mlp_mixer_3d_pytorch.py

@@ -13,13 +13,12 @@ def __init__(self, dim, fn):
     def forward(self, x):
         return self.fn(self.norm(x)) + x
 
-def FeedForward(dim, expansion_factor = 4, dropout = 0., dense = nn.Linear):
-    inner_dim = int(dim * expansion_factor)
+def FeedForward(dim, dim_hidden, dropout = 0., dense = nn.Linear):
     return nn.Sequential(
-        dense(dim, inner_dim),
+        dense(dim, dim_hidden),
         nn.GELU(),
         nn.Dropout(dropout),
-        dense(inner_dim, dim),
+        dense(dim_hidden, dim),
         nn.Dropout(dropout)
     )
 
@@ -35,8 +34,8 @@ def MLPMixer3D(*, image_size, time_size, channels, patch_size, time_patch_size,
         Rearrange('b c (t pt) (h p1) (w p2) -> b (h w t) (p1 p2 pt c)', p1 = patch_size, p2 = patch_size, pt = time_patch_size),
         nn.Linear((time_patch_size * patch_size ** 2) * channels, dim),
         *[nn.Sequential(
-            PreNormResidual(dim, FeedForward(num_patches, expansion_factor, dropout, chan_first)),
-            PreNormResidual(dim, FeedForward(dim, expansion_factor_token, dropout, chan_last))
+            PreNormResidual(dim, FeedForward(num_patches, int(expansion_factor * dim), dropout, chan_first)),
+            PreNormResidual(dim, FeedForward(dim, int(expansion_factor_token * dim), dropout, chan_last))
         ) for _ in range(depth)],
         nn.LayerNorm(dim),
         Reduce('b n c -> b c', 'mean'),

mlp_mixer_pytorch/mlp_mixer_pytorch.py

@@ -13,13 +13,12 @@ def __init__(self, dim, fn):
     def forward(self, x):
         return self.fn(self.norm(x)) + x
 
-def FeedForward(dim, expansion_factor = 4, dropout = 0., dense = nn.Linear):
-    inner_dim = int(dim * expansion_factor)
+def FeedForward(dim, dim_hidden, dropout = 0., dense = nn.Linear):
     return nn.Sequential(
-        dense(dim, inner_dim),
+        dense(dim, dim_hidden),
         nn.GELU(),
         nn.Dropout(dropout),
-        dense(inner_dim, dim),
+        dense(dim_hidden, dim),
         nn.Dropout(dropout)
     )
 
@@ -33,8 +32,8 @@ def MLPMixer(*, image_size, channels, patch_size, dim, depth, num_classes, expan
         Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
         nn.Linear((patch_size ** 2) * channels, dim),
         *[nn.Sequential(
-            PreNormResidual(dim, FeedForward(num_patches, expansion_factor, dropout, chan_first)),
-            PreNormResidual(dim, FeedForward(dim, expansion_factor_token, dropout, chan_last))
+            PreNormResidual(dim, FeedForward(num_patches, expansion_factor * dim, dropout, chan_first)),
+            PreNormResidual(dim, FeedForward(dim, expansion_factor_token * dim, dropout, chan_last))
         ) for _ in range(depth)],
         nn.LayerNorm(dim),
         Reduce('b n c -> b c', 'mean'),

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'mlp-mixer-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.0',
+  version = '0.3.0',
   license='MIT',
   description = 'MLP Mixer - Pytorch',
   long_description_content_type = 'text/markdown',