Create nets/encoder.py

nets/encoder.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from nets.multi_headed_attention import MultiHeadAttention
+import math
+
+
+class SkipConnection(nn.Module):
+    def __init__(self, module, use_mask=True):
+        super(SkipConnection, self).__init__()
+        self.use_mask = use_mask
+        self.module = module
+
+    def forward(self, input):
+        if isinstance(input, tuple):
+            if len(input) > 1:
+                input, mask = input[0], input[1]
+            else:
+                input = input[0]
+        else:
+            mask = None
+
+        if self.use_mask:
+            return input + self.module(input, mask=mask), mask
+        else:
+            return input + self.module(input), mask
+
+
+class Normalization(nn.Module):
+    def __init__(self, embed_dim, normalization='batch'):
+        super(Normalization, self).__init__()
+
+        normalizer_class = {
+            'batch': nn.BatchNorm1d,
+            'instance': nn.InstanceNorm1d
+        }.get(normalization, None)
+
+        self.normalizer = normalizer_class(embed_dim, affine=True)
+
+    def forward(self, input):
+        if isinstance(input, tuple):
+            if len(input) > 1:
+                input, mask = input[0], input[1]
+            else:
+                input = input[0]
+        else:
+            mask = None
+
+        if isinstance(self.normalizer, nn.BatchNorm1d):
+            return self.normalizer(input.view(-1, input.size(-1))).view(*input.size()), mask
+        elif isinstance(self.normalizer, nn.InstanceNorm1d):
+            return self.normalizer(input.permute(0, 2, 1)).permute(0, 2, 1), mask
+        else:
+            assert self.normalizer is None, "Unknown normalizer type"
+            return input, mask
+
+
+class MultiHeadAttentionLayer(nn.Sequential):
+    def __init__(self, n_heads, embed_dim, feed_forward_hidden=512, normalization='batch'):
+        super(MultiHeadAttentionLayer, self).__init__(
+            SkipConnection(
+                MultiHeadAttention(n_heads, input_dim=embed_dim, embed_dim=embed_dim),
+                use_mask=True
+            ),
+            Normalization(embed_dim, normalization),
+            SkipConnection(
+                nn.Sequential(
+                    nn.Linear(embed_dim, feed_forward_hidden),
+                    nn.ReLU(),
+                    nn.Linear(feed_forward_hidden, embed_dim)
+                ) if feed_forward_hidden > 0 else nn.Linear(embed_dim, embed_dim),
+                use_mask=False
+            ),
+            Normalization(embed_dim, normalization)
+        )
+
+
+class Encoder(nn.Module):
+    def __init__(self, n_heads, embed_dim, n_layers, node_dim=None,
+                 normalization='batch', feed_forward_hidden=200):
+        super(Encoder, self).__init__()
+
+        self.init_embed = nn.Linear(node_dim, embed_dim) if node_dim is not None else None
+
+        self.layers = nn.Sequential(*(
+            MultiHeadAttentionLayer(
+                n_heads, embed_dim,
+                feed_forward_hidden=feed_forward_hidden,
+                normalization=normalization
+            ) for _ in range(n_layers)
+        ))
+
+    def forward(self, input, mask=None):
+        device = input.device
+        batch_size = input.shape[0]
+        num_nodes = input.shape[1]
+
+        if mask is None:
+            mask = torch.ones(batch_size, num_nodes, num_nodes).to(device).float()
+
+        mask = (mask == 0)  # invert mask: 1s where we want to mask
+
+        x = input
+        h = self.init_embed(x.view(-1, x.size(-1))).view(*x.size()[:2], -1) if self.init_embed is not None else x
+        h, _ = self.layers((h, mask))  # Pass both h and mask through layers
+        return h