add the PEG module from vit literature

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -215,6 +215,22 @@ def forward(
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)
 
+class PEG(nn.Module):
+    def __init__(
+        self,
+        dim
+    ):
+        super().__init__()
+        self.ds_conv = nn.Conv2d(dim, dim, 3, padding = 1, groups = dim)
+
+    def forward(self, x):
+        b, n, d = x.shape
+        hw = int(math.sqrt(n))
+        x = rearrange(x, 'b (h w) d -> b d h w', h = hw)
+        x = self.ds_conv(x)
+        x = rearrange(x, 'b d h w -> b (h w) d')
+        return x
+
 class FeedForward(nn.Module):
     def __init__(self, dim, mult = 4, time_cond_dim = None):
         super().__init__()
@@ -326,6 +342,7 @@ def __init__(
                 FeedForward(dim)
             ]))
 
+        self.patches_peg = PEG(dim)
         self.patches_self_attn = LinearAttention(dim, norm = True, **attn_kwargs)
         self.patches_self_attn_ff = FeedForward(dim)
 
@@ -370,6 +387,8 @@ def forward(
         # the recurrent interface network body
 
         for _ in range(self.depth):
+            patches = self.patches_peg(patches) + patches
+
             # latents extract or cluster information from the patches
 
             latents = self.latents_attend_to_patches(latents, patches, time = t) + latents

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'RIN-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.3',
+  version = '0.0.5',
   license='MIT',
   description = 'RIN - Recurrent Interface Network - Pytorch',
   author = 'Phil Wang',