From 6e3ef21b1ff65b9e6570fbe50a9f037d2b33e963 Mon Sep 17 00:00:00 2001
From: Piotr Czapla <github@piotrczapla.com>
Date: Wed, 21 Nov 2018 18:44:49 +0100
Subject: [PATCH] Add Avg BiClassifier

---
 fastai_contrib/learner.py |  2 +-
 fastai_contrib/models.py  | 39 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/fastai_contrib/learner.py b/fastai_contrib/learner.py
index f3dc5dc..a1dfc27 100644
--- a/fastai_contrib/learner.py
+++ b/fastai_contrib/learner.py
@@ -65,7 +65,7 @@ def bilm_text_classifier_learner(data: DataBunch, bptt: int = 70, max_len: int =
     if ps is None:  ps = [0.1]
     ds = data.train_ds
     vocab_size, n_class = len(data.vocab.itos), data.c
-    layers = [emb_sz * 3 * 2] + lin_ftrs + [n_class]
+    layers = [emb_sz * 3] + lin_ftrs + [n_class]
     ps = [dps[4]] + ps
     model = get_birnn_classifier(bptt, max_len, n_class, vocab_size, emb_sz, nh, nl, pad_token,
                                layers, ps, input_p=dps[0], weight_p=dps[1], embed_p=dps[2], hidden_p=dps[3],
diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py
index 3d04bdb..b99db97 100644
--- a/fastai_contrib/models.py
+++ b/fastai_contrib/models.py
@@ -79,6 +79,43 @@ class BiPoolingLinearClassifier(nn.Module):
             x = self.layers(x)
             return x, raw_outputs, outputs
 
+
+class AvgPoolingLinearClassifier(nn.Module):
+    "Create a linear classifier with pooling."
+
+    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.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]
+        if len(output.size()) == 3:
+            sl,bs,_ = output.size()
+            avgpool = self.pool(output, bs, False)
+            mxpool = self.pool(output, bs, True)
+            x = torch.cat([output[-1], mxpool, avgpool], 1)
+            x = self.layers(x)
+            return x, raw_outputs, outputs
+        elif len(output.size()) == 4:
+            sl, bs, em_sz, passes = output.size()
+
+            avgpool = (self.pool(output[...,0], bs, False) + self.pool(output[..., 1], bs, False))/2
+            mxpool = (self.pool(output[...,0], bs, True) +self.pool(output[..., 1], bs, True))/2
+            x = torch.cat([(output[-1][..., 0]+output[-1][..., 1])/2, mxpool, avgpool], 1)
+            x = self.layers(x)
+            return x, raw_outputs, outputs
+
+
 def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, tie_weights:bool=True,
                        qrnn:bool=False, bias:bool=True, bidir:bool=False, output_p:float=0.4, hidden_p:float=0.2, input_p:float=0.6,
                        embed_p:float=0.1, weight_p:float=0.5)->nn.Module:
@@ -106,6 +143,6 @@ def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_s
     bwd_rnn_enc = MultiBatchRNNCore(bptt, max_seq, vocab_sz, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                                 qrnn=qrnn, hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
 
-    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), BiPoolingLinearClassifier(layers, drops))
+    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), AvgPoolingLinearClassifier(layers, drops))
     model.reset()
     return model
\ No newline at end of file