add lambda convolutions, as described in the paper

Author: lucidrains

README.md

@@ -10,6 +10,8 @@ $ pip install lambda-networks
 
 ## Usage
 
+Global context
+
 ```python
 import torch
 from lambda_networks import LambdaLayer
@@ -18,7 +20,6 @@ layer = LambdaLayer(
     dim = 32,       # channels going in
     dim_out = 32,   # channels out
     n = 64 * 64,    # number of input pixels (64 x 64 image)
-    m = 64 * 64,    # number of context (64 x 64 global)
     dim_k = 16,     # key dimension
     heads = 4,      # number of heads, for multi-query
     dim_u = 1       # 'intra-depth' dimension
@@ -28,9 +29,28 @@ x = torch.randn(1, 32, 64, 64)
 layer(x) # (1, 32, 64, 64)
 ```
 
+Localized context
+
+```python
+import torch
+from lambda_networks import LambdaLayer
+
+layer = LambdaLayer(
+    dim = 32,
+    dim_out = 32,
+    r = 23,         # the receptive field for relative positional encoding (23 x 23)
+    dim_k = 16,
+    heads = 4,
+    dim_u = 4
+)
+
+x = torch.randn(1, 32, 64, 64)
+layer(x) # (1, 32, 64, 64)
+```
+
 ## Todo
 
-- [ ] Lambda layers with structured context
+- [x] Lambda layers with structured context
 - [ ] Document hyperparameters and put some sensible defaults
 - [ ] Test it out
 

lambda_networks/lambda_networks.py

@@ -18,9 +18,9 @@ def __init__(
         self,
         dim,
         *,
-        n,
-        m,
         dim_k,
+        n = None,
+        r = None,
         heads = 4,
         dim_out = None,
         dim_u = 1):
@@ -35,11 +35,20 @@ def __init__(
         self.to_q = nn.Conv2d(dim, dim_k * heads, 1, bias = False)
         self.to_k = nn.Conv2d(dim, dim_k * dim_u, 1, bias = False)
         self.to_v = nn.Conv2d(dim, dim_v * dim_u, 1, bias = False)
-        self.pos_emb = nn.Parameter(torch.randn(n, m, dim_k, dim_u))
 
         self.norm_q = nn.BatchNorm2d(dim_k * heads)
         self.norm_v = nn.BatchNorm2d(dim_v * dim_u)
 
+        self.local_contexts = exists(r)
+        if exists(r):
+            assert (r % 2) == 1, 'Receptive kernel size should be odd'
+            self.padding = r // 2
+            self.R = nn.Parameter(torch.randn(dim_k, dim_u, 1, r, r))
+        else:
+            assert exists(n), 'You must specify the total sequence length (h x w)'
+            self.pos_emb = nn.Parameter(torch.randn(n, n, dim_k, dim_u))
+
+
     def forward(self, x):
         b, c, hh, ww, u, h = *x.shape, self.u, self.heads
 
@@ -57,10 +66,15 @@ def forward(self, x):
         k = k.softmax(dim=-1)
 
         λc = einsum('b k u m, b v u m -> b k v', k, v)
-        λp = einsum('n m k u, b v u m -> b n k v', self.pos_emb, v)
-
         Yc = einsum('b h k n, b k v -> b n h v', q, λc)
-        Yp = einsum('b h k n, b n k v -> b n h v', q, λp)
+
+        if self.local_contexts:
+            v = rearrange(v, 'b v u (hh ww) -> b u v hh ww', hh = hh, ww = ww)
+            λp = F.conv3d(v, self.R, padding = (0, self.padding, self.padding))
+            Yp = einsum('b h k n, b k v n -> b n h v', q, λp.flatten(3))
+        else:
+            λp = einsum('n m k u, b v u m -> b n k v', self.pos_emb, v)
+            Yp = einsum('b h k n, b n k v -> b n h v', q, λp)
 
         Y = Yc + Yp
         out = rearrange(Y, 'b (hh ww) h v -> b (h v) hh ww', hh = hh, ww = ww)

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'lambda-networks',
   packages = find_packages(),
-  version = '0.0.2',
+  version = '0.1.0',
   license='MIT',
   description = 'Lambda Networks - Pytorch',
   author = 'Phil Wang',