add token shift feature, which should greatly improve convergence. bump to 1.0

Author: lucidrains

README.md

@@ -147,8 +147,9 @@ images.shape # (4, 3, 256, 256)
 ```
 
 You may also want to generate text using DALL-E. For that call this function:
-```
-text_tokens, texts = dalle.generate_texts(text)
+
+```python
+text_tokens, texts = dalle.generate_texts(tokenizer, text)
 ```
 
 ## OpenAI's Pretrained VAE

dalle_pytorch/dalle_pytorch.py

@@ -7,7 +7,7 @@
 from axial_positional_embedding import AxialPositionalEmbedding
 from einops import rearrange
 
-from dalle_pytorch import distributed_utils, tokenizer
+from dalle_pytorch import distributed_utils
 from dalle_pytorch.vae import OpenAIDiscreteVAE, VQGanVAE
 from dalle_pytorch.transformer import Transformer, DivideMax
 
@@ -322,7 +322,8 @@ def __init__(
         sparse_attn = False,
         attn_types = None,
         loss_img_weight = 7,
-        stable = False
+        stable = False,
+        shift_tokens = True
     ):
         super().__init__()
         assert isinstance(vae, (DiscreteVAE, OpenAIDiscreteVAE, VQGanVAE)), 'vae must be an instance of DiscreteVAE'
@@ -367,7 +368,8 @@ def __init__(
             attn_types = attn_types,
             image_fmap_size = image_fmap_size,
             sparse_attn = sparse_attn,
-            stable = stable
+            stable = stable,
+            shift_tokens = shift_tokens
         )
 
         self.stable = stable
@@ -399,7 +401,8 @@ def __init__(
     @eval_decorator
     def generate_texts(
         self,
-        text=None,
+        tokenizer,
+        text = None,
         *,
         filter_thres = 0.5,
         temperature = 1.
@@ -577,5 +580,3 @@ def forward(
 
         loss = (loss_text + self.loss_img_weight * loss_img) / (self.loss_img_weight + 1)
         return loss
-    
-    

dalle_pytorch/transformer.py

@@ -52,6 +52,8 @@ def __init__(self, dim, depth, fn):
     def forward(self, x, **kwargs):
         return self.fn(x, **kwargs) * self.scale
 
+# layer norm
+
 class PreNorm(nn.Module):
     def __init__(self, dim, fn):
         super().__init__()
@@ -61,6 +63,8 @@ def __init__(self, dim, fn):
     def forward(self, x, **kwargs):
         return self.fn(self.norm(x), **kwargs)
 
+# feed forward
+
 class GEGLU(nn.Module):
     def forward(self, x):
         x, gates = x.chunk(2, dim = -1)
@@ -79,6 +83,49 @@ def __init__(self, dim, dropout = 0., mult = 4.):
     def forward(self, x):
         return self.net(x)
 
+# token shift classes
+
+class PreShiftToken(nn.Module):
+    def __init__(self, fn, image_size, seq_len):
+        super().__init__()
+        self.fn = fn
+        self.image_size = image_size
+        self.seq_len = seq_len
+
+    def forward(self, x, **kwargs):
+        n = x.shape[1]
+        seq_len, image_size = self.seq_len, self.image_size
+        img_seq_len = image_size ** 2
+        text_len = seq_len - img_seq_len + 1
+        padding = seq_len - n + 1
+
+        # get text and image tokens
+
+        x_text, x_img = x[:, :text_len], x[:, text_len:]
+        x_img = F.pad(x_img, (0, 0, 0, padding))
+        x_img = rearrange(x_img, 'b (h w) d -> b h w d', h = image_size)
+
+        # shift 1 from the left for text tokens
+
+        x_text_shift, x_text_pass = x_text.chunk(2, dim = -1)
+        x_text_shift = F.pad(x_text_shift, (0, 0, 1, -1))
+        x_text = torch.cat((x_text_shift, x_text_pass), dim = -1)
+
+        # shift from top, left for image tokens
+
+        x_img_shift_top, x_img_shift_left, *x_img_pass = x_img.chunk(4, dim = -1)
+        x_img_shift_left = F.pad(x_img_shift_left, (0, 0, 1, -1))
+        x_img_shift_top = F.pad(x_img_shift_top, (0, 0, 0, 0, 1, -1))
+        x_img = torch.cat((x_img_shift_top, x_img_shift_left, *x_img_pass), dim = -1)
+
+        # merge text and image sequence back together
+
+        x_img = rearrange(x_img, 'b h w d -> b (h w) d')
+        x = torch.cat((x_text, x_img[:, :-padding]), dim = 1)
+        return self.fn(x, **kwargs)
+
+# main transformer class
+
 class Transformer(nn.Module):
     def __init__(
         self,
@@ -96,7 +143,8 @@ def __init__(
         attn_types = None,
         image_fmap_size = None,
         sparse_attn = False,
-        stable = False
+        stable = False,
+        shift_tokens = True
     ):
         super().__init__()
         layers = nn.ModuleList([])
@@ -127,9 +175,14 @@ def __init__(
             else:
                 attn = attn_class(dim = dim, causal = causal, dim_ff = dim * 4)
 
+            ff = FeedForward(dim, mult = ff_mult, dropout = ff_dropout)
+
+            if shift_tokens:
+                attn, ff = map(lambda t: PreShiftToken(t, image_size = image_fmap_size, seq_len = seq_len), (attn, ff))
+
             layers.append(nn.ModuleList([
                 LayerScale(dim, ind + 1, PreNorm(dim, attn)),
-                LayerScale(dim, ind + 1, PreNorm(dim, FeedForward(dim, mult = ff_mult, dropout = ff_dropout)))
+                LayerScale(dim, ind + 1, PreNorm(dim, ff))
             ]))
 
         execute_type = ReversibleSequence if reversible else SequentialSequence

generate.py

@@ -103,7 +103,7 @@ def exists(val):
 
 for j, text in tqdm(enumerate(texts)):
     if args.gentxt:
-        text_tokens, gen_texts = dalle.generate_texts(text=text, filter_thres = args.top_k)
+        text_tokens, gen_texts = dalle.generate_texts(tokenizer, text=text, filter_thres = args.top_k)
         text = gen_texts[0]
     else:
         text_tokens = tokenizer.tokenize([text], dalle.text_seq_len).cuda()

setup.py

@@ -4,7 +4,7 @@
   name = 'dalle-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '0.14.3',
+  version = '1.0.0',
   license='MIT',
   description = 'DALL-E - Pytorch',
   author = 'Phil Wang',