From d0c15472d6b6e7cae441ad9b31c9dc9dbf3438fb Mon Sep 17 00:00:00 2001
From: Aayush <aayushyadav96@gmail.com>
Date: Sun, 6 Jan 2019 23:58:09 +0530
Subject: [PATCH] Self attention for pooling linear classifier

This PR will introduce a `BiAttentionPoolingClassifier` as in [Attention is all you need](https://arxiv.org/abs/1706.03762) following the discussion with @sebastianruder in Teams.

I ran out of memory on my 1060 while testing the attention module, but was able to at least verify that it is functionally correct. Some changes might be required to ensure that the tensor passed to `self.layers` is of the right shape (but I'm not quite sure as of now).

I'll shift all the stuff to Collab for testing and see if it's any help.
---
 fastai_contrib/models.py | 116 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 115 insertions(+), 1 deletion(-)

diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py
index 92ac7e4..61a84c1 100644
--- a/fastai_contrib/models.py
+++ b/fastai_contrib/models.py
@@ -70,6 +70,120 @@ class MultiBatchBiLMModel(BiLMModel):
                 outputs.append(o)
         return self.concat(raw_outputs), self.concat(outputs)
 
+class BiAttentionPoolingClassifier(nn.Module):
+    r" [WIP] BiLM Pooling with self attention"
+
+    def __init__(self, layers:Collection[int], drops:Collection[float]):
+        super().__init__()
+        mod_layers = []
+        activs = [nn.ReLU(inplace=True)] * (len(layers) - 2) + [None]
+        for n_in,n_out,p,actn in zip(layers[:-1],layers[1:], drops, activs):
+            mod_layers += bn_drop_lin(n_in, n_out, p=p, actn=actn)
+        self.self_attn = MultiHeadAttention(n_head=8, d_model=1, d_k=64, d_v=64, dropout=0.1)
+        self.layers = nn.Sequential(*mod_layers)
+
+    def pool(self, x:Tensor, bs:int, is_max:bool):
+        "Pool the tensor along the seq_len dimension."
+        f = F.adaptive_max_pool1d if is_max else F.adaptive_avg_pool1d
+        return f(x.permute(1,2,0), (1,)).view(bs,-1)
+    
+    def forward(self, input:Tuple[Tensor,Tensor])->Tuple[Tensor,Tensor,Tensor]:
+        raw_outputs, outputs = input
+        output = outputs[-1]
+        assert len(output.size()) == 4, 'Expected input dimension 4'
+        sl, bs, em_sz, passes = output.size()
+
+        f_avgpool = self.pool(output[...,0], bs, False)
+        f_mxpool = self.pool(output[...,0], bs, True)
+        b_avgpool = self.pool(output[..., 1], bs, False)
+        b_mxpool = self.pool(output[..., 1], bs, True)
+        x = torch.cat([output[-1][..., 0], f_mxpool, f_avgpool,
+                        output[-1][..., 1], b_mxpool, b_avgpool,], 1)
+        x = x.unsqueeze(-1)
+        x, _ = self.self_attn(x, x, x)
+        
+        x = self.layers(x)
+        return x, raw_outputs, outputs
+
+class ScaledDotProductAttention(nn.Module):
+    r""" 
+    Scaled Dot-Product Attention
+    based on: https://github.com/jadore801120/attention-is-all-you-need-pytorch
+    """
+
+    def __init__(self, temperature, attn_dropout=0.1):
+        super().__init__()
+        self.temperature = temperature
+        self.dropout = nn.Dropout(attn_dropout)
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, q, k, v):
+
+        attn = torch.bmm(q, k.transpose(1, 2))
+        attn = attn / self.temperature
+
+        attn = self.softmax(attn)
+        attn = self.dropout(attn)
+        output = torch.bmm(attn, v)
+
+        return output, attn
+
+class MultiHeadAttention(nn.Module):
+    r""" 
+    Multi-Head Attention module
+    based on: https://github.com/jadore801120/attention-is-all-you-need-pytorch
+    """
+
+    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
+        super().__init__()
+
+        self.n_head = n_head
+        self.d_k = d_k
+        self.d_v = d_v
+
+        self.w_qs = nn.Linear(d_model, n_head * d_k)
+        self.w_ks = nn.Linear(d_model, n_head * d_k)
+        self.w_vs = nn.Linear(d_model, n_head * d_v)
+        nn.init.normal_(self.w_qs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
+        nn.init.normal_(self.w_ks.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_k)))
+        nn.init.normal_(self.w_vs.weight, mean=0, std=np.sqrt(2.0 / (d_model + d_v)))
+
+        self.attention = ScaledDotProductAttention(temperature=np.power(d_k, 0.5))
+        self.layer_norm = nn.LayerNorm(d_model)
+
+        self.fc = nn.Linear(n_head * d_v, d_model)
+        nn.init.xavier_normal_(self.fc.weight)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, q, k, v):
+
+        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
+
+        
+        sz_b, len_q, _ = q.size()
+        sz_b, len_k, _ = k.size()
+        sz_b, len_v, _ = v.size()
+        residual = q
+
+        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
+        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
+        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
+
+        q = q.permute(2, 0, 1, 3).contiguous().view(-1, len_q, d_k) # (n*b) x lq x dk
+        k = k.permute(2, 0, 1, 3).contiguous().view(-1, len_k, d_k) # (n*b) x lk x dk
+        v = v.permute(2, 0, 1, 3).contiguous().view(-1, len_v, d_v) # (n*b) x lv x dv
+
+        x, attn = self.attention(q, k, v)
+
+        x = x.view(n_head, sz_b, len_q, d_v)
+        x = x.permute(1, 2, 0, 3).contiguous().view(sz_b, len_q, -1) # b x lq x (n*dv)
+
+        x = self.dropout(self.fc(x))
+        x = self.layer_norm(x + residual)
+
+        return x, attn    
+
 class BiPoolingLinearClassifier(PoolingLinearClassifier):
     "Create a linear classifier with pooling."
 
@@ -163,4 +277,4 @@ def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_s
     model.reset()
     return model
 
-#endregion
\ No newline at end of file
+#endregion