token shift within feedforward best accompanied by a layernorm, from …

…normformer paper
lucidrains · Mar 19, 2023 · 3280b3d · 3280b3d
1 parent bde9f22
commit 3280b3d
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -176,3 +176,14 @@ video_as_images_out = unet(video, enable_time = False)
     year    = {2021}
 }
 ```
+
+```bibtex
+@inproceedings{anonymous2022normformer,
+    title   = {NormFormer: Improved Transformer Pretraining with Extra Normalization},
+    author  = {Anonymous},
+    booktitle = {Submitted to The Tenth International Conference on Learning Representations },
+    year    = {2022},
+    url     = {https://openreview.net/forum?id=GMYWzWztDx5},
+    note    = {under review}
+}
+```
diff --git a/make_a_video_pytorch/make_a_video.py b/make_a_video_pytorch/make_a_video.py
@@ -78,7 +78,10 @@ def __init__(self, dim, mult = 4):
             GEGLU()
         )
 
-        self.proj_out = nn.Conv3d(inner_dim, dim, 1, bias = False)
+        self.proj_out = nn.Sequential(
+            LayerNorm(inner_dim),
+            nn.Conv3d(inner_dim, dim, 1, bias = False)
+        )
 
     def forward(self, x, enable_time = True):
         x = self.proj_in(x)

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'make-a-video-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.1',
+  version = '0.1.2',
   license='MIT',
   description = 'Make-A-Video - Pytorch',
   author = 'Phil Wang',