Modify attention implementation

1 files changed, 68 insertions, 12 deletions

diff --git a/model/attention.py b/model/attention.py
index d82ca17..48c9e29 100644
--- a/model/attention.py
+++ b/model/attention.py
@@ -1,15 +1,71 @@
+import numpy as np
+import torch
 from torch import nn
 
+class Projection(nn.Module):
+    """
+    Linear transformation by a matrix.
+    """
+    def __init__(self, in_dim, out_dim):
+        super(Projection, self).__init__()
+        self.projection = nn.Parameter(torch.normal(torch.zeros(in_dim, out_dim), 1, ))
+
+    def forward(self, x):
+        """
+        x shape: batch, seqlen, in_dim
+        """
+        return x.matmul(self.projection)
+
+class Head(nn.Module):
+    def __init__(self, hidden_dim, output_dim):
+        super(Head, self).__init__()
+        self._scaling_factor = np.sqrt(output_dim)
+        self.query_projection = Projection(hidden_dim, output_dim)
+        self.key_projection = Projection(hidden_dim, output_dim)
+        self.value_projection = Projection(hidden_dim, output_dim)
+
+    def forward(self, x):
+        """
+        x shape: seqlen, batch, hiddendim
+        """
+        # get batch in front
+        x = torch.transpose(x, 0, 1)
+
+        query = self.query_projection(x)
+        key = self.key_projection(x)
+        value = self.value_projection(x)
+
+        # transpose the matrix dimensions of key to align for multiplication
+        product = query.matmul(torch.transpose(key, 1, 2))
+        weights = torch.softmax(product / self._scaling_factor, dim=2)
+        value = weights.matmul(value)
+        value = torch.transpose(value, 0, 1)
+        weights = torch.transpose(weights, 0, 1)
+
+        return value, weights
+
 class Attention(nn.Module):
-  def __init__(self, hidden_dim, num_heads):
-    super(Attention, self).__init__()
-    # TODO: implement Attention
-    pass
-
-  def forward(self, x):
-    # TODO: implement Attention; return both result of attention mechanism and
-    # attention weights (for visualization).
-    # x shape: (seqlen, batch, hiddendim)
-    result, att_weights = x, None # placeholder
-    pass
-    return result, att_weights
+    def __init__(self, hidden_dim, num_heads):
+        super(Attention, self).__init__()
+        self._num_heads = num_heads
+        self._head_output_dim = hidden_dim // num_heads
+        # ensure hidden_dim is divisible by num_heads
+        assert(self._head_output_dim * num_heads == hidden_dim)
+        self.heads = nn.ModuleList([
+            Head(hidden_dim, self._head_output_dim) for _ in range(num_heads)
+        ])
+        self.final_projection = nn.Linear(hidden_dim, hidden_dim)
+        pass
+
+    def forward(self, x):
+        # x shape: (seqlen, batch, hiddendim)
+        result = torch.zeros(x.shape)
+        # attentions are (heads, seqlen, batch, seqlen)
+        attentions = torch.zeros(self._num_heads, x.shape[0], x.shape[1], x.shape[0])
+        for i in range(self._num_heads):
+            from_index = i * self._head_output_dim
+            to_index = from_index + self._head_output_dim
+            result[:, :, from_index:to_index], attentions[i, :, :, :] = self.heads[i](x)
+
+        result = self.final_projection(result)
+        return result, attentions