seeing a signal with dual patchnorm in another repository, fully inco…

…rporate
isbee · Feb 6, 2023 · 46dcaf2 · 46dcaf2
1 parent bdaf2d1
commit 46dcaf2
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '1.0.0',
+  version = '1.0.1',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   long_description_content_type = 'text/markdown',

diff --git a/vit_pytorch/learnable_memory_vit.py b/vit_pytorch/learnable_memory_vit.py
@@ -118,7 +118,9 @@ def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, ml
 
         self.to_patch_embedding = nn.Sequential(
             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim)
         )
 
         self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

diff --git a/vit_pytorch/twins_svt.py b/vit_pytorch/twins_svt.py
@@ -71,7 +71,12 @@ def __init__(self, *, dim, dim_out, patch_size):
         self.dim = dim
         self.dim_out = dim_out
         self.patch_size = patch_size
-        self.proj = nn.Conv2d(patch_size ** 2 * dim, dim_out, 1)
+
+        self.proj = nn.Sequential(
+            LayerNorm(patch_size ** 2 * dim),
+            nn.Conv2d(patch_size ** 2 * dim, dim_out, 1),
+            LayerNorm(dim_out)
+        )
 
     def forward(self, fmap):
         p = self.patch_size

diff --git a/vit_pytorch/vit_with_patch_merger.py b/vit_pytorch/vit_with_patch_merger.py
@@ -121,7 +121,9 @@ def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, ml
 
         self.to_patch_embedding = nn.Sequential(
             Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
             nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim)
         )
 
         self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))