add a feedforward after the self-attention

lucidrains · Jun 29, 2020 · 8a24fb6 · 8a24fb6
1 parent e40eb6f
commit 8a24fb6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'stylegan2_pytorch',
   packages = find_packages(),
   scripts=['bin/stylegan2_pytorch'],
-  version = '0.17.2',
+  version = '0.17.3',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

diff --git a/stylegan2_pytorch/stylegan2_pytorch.py b/stylegan2_pytorch/stylegan2_pytorch.py
@@ -87,6 +87,13 @@ def forward(self, x):
         out = out.permute(0, 3, 1, 2)
         return out, loss
 
+# one layer of self-attention and feedforward, for images
+
+attn_and_ff = lambda chan: nn.Sequential(*[
+    Residual(Rezero(ImageLinearAttention(chan))),
+    Residual(Rezero(nn.Sequential(nn.Conv2d(chan, chan * 2, 1), leaky_relu(), nn.Conv2d(chan * 2, chan, 1))))
+])
+
 # helpers
 
 def default(value, d):
@@ -424,9 +431,7 @@ def __init__(self, image_size, latent_dim, network_capacity = 16, transparent =
             not_last = ind != (self.num_layers - 1)
             num_layer = self.num_layers - ind
 
-            attn_fn = nn.Sequential(*[
-                Residual(Rezero(ImageLinearAttention(in_chan))) for _ in range(2)
-            ]) if num_layer in attn_layers else None
+            attn_fn = attn_and_ff(in_chan) if num_layer in attn_layers else None
 
             self.attns.append(attn_fn)
 
@@ -485,7 +490,8 @@ def __init__(self, image_size, network_capacity = 16, fq_layers = [], fq_dict_si
             blocks.append(block)
 
             attn_fn = nn.Sequential(*[
-                Residual(Rezero(ImageLinearAttention(out_chan))) for _ in range(2)
+                Residual(Rezero(ImageLinearAttention(out_chan))),
+                Residual(Rezero(nn.Sequential(nn.Conv2d(out_chan, out_chan, 1), leaky_relu(), nn.Conv2d(out_chan, out_chan, 1))))
             ]) if num_layer in attn_layers else None
 
             attn_blocks.append(attn_fn)