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

@@ -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
+#endregion