Update nets/encoder.py

nets/encoder.py

@@ -17,6 +17,7 @@ class SkipConnection(nn.Module):
                 input, mask = input[0], input[1]
             else:
                 input = input[0]
+                mask = None
         else:
             mask = None
 
@@ -29,7 +30,6 @@ class SkipConnection(nn.Module):
 class Normalization(nn.Module):
     def __init__(self, embed_dim, normalization='batch'):
         super(Normalization, self).__init__()
-
         normalizer_class = {
             'batch': nn.BatchNorm1d,
             'instance': nn.InstanceNorm1d
@@ -43,6 +43,7 @@ class Normalization(nn.Module):
                 input, mask = input[0], input[1]
             else:
                 input = input[0]
+                mask = None
         else:
             mask = None
 
@@ -51,28 +52,33 @@ class Normalization(nn.Module):
         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):
+class MultiHeadAttentionLayer(nn.Module):
     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)
+        super(MultiHeadAttentionLayer, self).__init__()
+        self.attention = SkipConnection(
+            MultiHeadAttention(n_heads, input_dim=embed_dim, embed_dim=embed_dim),
+            use_mask=True
+        )
+        self.norm1 = Normalization(embed_dim, normalization)
+        self.ff = 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
         )
+        self.norm2 = Normalization(embed_dim, normalization)
+
+    def forward(self, input):
+        h, mask = self.attention(input)
+        h, mask = self.norm1((h, mask))
+        h, mask = self.ff((h, mask))
+        h, mask = self.norm2((h, mask))
+        return h, mask
 
 
 class Encoder(nn.Module):
@@ -82,13 +88,13 @@ class Encoder(nn.Module):
 
         self.init_embed = nn.Linear(node_dim, embed_dim) if node_dim is not None else None
 
-        self.layers = nn.Sequential(*(
+        self.layers = nn.ModuleList([
             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
@@ -97,10 +103,11 @@ class Encoder(nn.Module):
 
         if mask is None:
             mask = torch.ones(batch_size, num_nodes, num_nodes).to(device).float()
+        mask = (mask == 0)
 
-        mask = (mask == 0)  # invert mask: 1s where we want to mask
+        x = self.init_embed(input) if self.init_embed is not None else input
+        h = x
+        for layer in self.layers:
+            h, mask = layer((h, 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