simplify mlp head

Update README and documentation

README.md

@@ -1,4 +1,4 @@
-<img src="./vit.png" width="500px"></img>
+<img src="./vit.gif" width="500px"></img>
 
 ## Vision Transformer - Pytorch
 
@@ -36,6 +36,28 @@ mask = torch.ones(1, 8, 8).bool() # optional mask, designating which patch to at
 preds = v(img, mask = mask) # (1, 1000)
 ```
 
+## Parameters
+- `image_size`: int.  
+Image size.
+- `patch_size`: int.  
+Number of patches. `image_size` must be divisible by `patch_size`.  
+The number of patches is: ` n = (image_size // patch_size) ** 2` and `n` **must be greater than 16**.
+- `num_classes`: int.  
+Number of classes to classify.
+- `dim`: int.  
+Last dimension of output tensor after linear transformation `nn.Linear(..., dim)`.
+- `depth`: int.  
+Number of Transformer blocks.
+- `heads`: int.  
+Number of heads in Multi-head Attention layer. 
+- `mlp_dim`: int.  
+Dimension of the MLP (FeedForward) layer. 
+- `channels`: int, default `3`.  
+Number of image's channels. 
+- `dropout`: float between `[0, 1]`, default `0.`.  
+Dropout rate. 
+- `emb_dropout`: float between `[0, 1]`, default `0`.  
+Embedding dropout rate.
 ## Research Ideas
 
 ### Self Supervised Training

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '0.2.7',
+  version = '0.3.0',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   author = 'Phil Wang',

vit_pytorch/efficient.py

@@ -20,9 +20,7 @@ class ViT(nn.Module):
 
         self.mlp_head = nn.Sequential(
             nn.LayerNorm(dim),
-            nn.Linear(dim, dim * 4),
-            nn.GELU(),
-            nn.Linear(dim * 4, num_classes)
+            nn.Linear(dim, num_classes)
         )
 
     def forward(self, img):

vit_pytorch/vit_pytorch.py

@@ -85,10 +85,10 @@ class Transformer(nn.Module):
 class ViT(nn.Module):
     def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels = 3, dropout = 0., emb_dropout = 0.):
         super().__init__()
-        assert image_size % patch_size == 0, 'image dimensions must be divisible by the patch size'
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
         num_patches = (image_size // patch_size) ** 2
         patch_dim = channels * patch_size ** 2
-        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective. try decreasing your patch size'
+        assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size'
 
         self.patch_size = patch_size
 
@@ -103,10 +103,7 @@ class ViT(nn.Module):
 
         self.mlp_head = nn.Sequential(
             nn.LayerNorm(dim),
-            nn.Linear(dim, mlp_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(mlp_dim, num_classes)
+            nn.Linear(dim, num_classes)
         )
 
     def forward(self, img, mask = None):