Working version of biclassfier

fastai_contrib/learner.py

@@ -27,6 +27,34 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in
     return learn
 
 
+
+def convert_weights(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:Collection[str]) -> Weights:
+    "Convert the model weights to go with a new vocabulary."
+    print(wgts.keys())
+    if 'fwd_lm.0.encoder.weight' in wgts: #todo share embedding matrix computation
+        wgts = convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='fwd_lm.')
+        return convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='bwd_lm.')
+    else:
+        return convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='')
+
+def convert_weights_with_prefix(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:Collection[str], prefix='') -> Weights:
+    "Convert the model weights to go with a new vocabulary."
+    dec_bias, enc_wgts = wgts[prefix+'1.decoder.bias'], wgts[prefix+'0.encoder.weight']
+    bias_m, wgts_m = dec_bias.mean(0), enc_wgts.mean(0)
+    new_w = enc_wgts.new_zeros((len(itos_new),enc_wgts.size(1))).zero_()
+    new_b = dec_bias.new_zeros((len(itos_new),)).zero_()
+    for i,w in enumerate(itos_new):
+        r = stoi_wgts[w] if w in stoi_wgts else -1
+        new_w[i] = enc_wgts[r] if r>=0 else wgts_m
+        new_b[i] = dec_bias[r] if r>=0 else bias_m
+    wgts[prefix+'0.encoder.weight'] = new_w
+    wgts[prefix+'0.encoder_dp.emb.weight'] = new_w.clone()
+    wgts[prefix+'1.decoder.weight'] = new_w.clone()
+    wgts[prefix+'1.decoder.bias'] = new_b
+    return wgts
+
+
+
 def bilm_text_classifier_learner(data: DataBunch, bptt: int = 70, max_len: int = 70 * 20, emb_sz: int = 400,
                             nh: int = 1150, nl: int = 3,
                             lin_ftrs: Collection[int] = None, ps: Collection[float] = None, pad_token: int = 1,
@@ -37,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] + lin_ftrs + [n_class]
+    layers = [emb_sz * 3 * 2] + 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],
@@ -59,56 +87,15 @@ def birnn_classifier_split(model:nn.Module) -> List[nn.Module]:
     groups.append([model[1]])
     return groups
 
-# learner extensions
-class RNNLearner(Learner):
-    "Basic class for a Learner in RNN."
-    def __init__(self, data:DataBunch, model:nn.Module, bptt:int=70, split_func:OptSplitFunc=None, clip:float=None,
-                 adjust:bool=False, alpha:float=2., beta:float=1., **kwargs):
-        super().__init__(data, model, **kwargs)
-        self.callbacks.append(RNNTrainer(self, bptt, alpha=alpha, beta=beta, adjust=adjust))
-        if clip: self.callback_fns.append(partial(GradientClipping, clip=clip))
-        if split_func: self.split(split_func)
-        self.metrics = [accuracy]
-
-    def model_path(self, name:str):
-        return self.path/self.model_dir/f'{name}.pth'
-
-    def _get_encoder(self):
-        return self.model.encoder if hasattr(self.model, 'encoder') else self.model[0]
-
-    def save_encoder(self, name:str):
-        "Save the encoder to `name` inside the model directory."
-        torch.save(self._get_encoder().state_dict(), self.model_path(name))
-
-    def load_encoder(self, name:str):
-        "Load the encoder `name` from the model directory."
-        self._get_encoder().load_state_dict(torch.load(self.model_path(name)))
-        self.freeze()
-
-    def load_pretrained(self, wgts_fname:str, itos_fname:str):
-        "Load a pretrained model and adapts it to the data vocabulary."
-        old_itos = pickle.load(open(itos_fname, 'rb'))
-        old_stoi = {v:k for k,v in enumerate(old_itos)}
-        wgts = torch.load(wgts_fname, map_location=lambda storage, loc: storage)
-        wgts = convert_weights(wgts, old_stoi, self.data.train_ds.vocab.itos)
-        self.model.load_state_dict(wgts)
-
-    def get_preds(self, ds_type:DatasetType=DatasetType.Valid, with_loss:bool=False, n_batch:Optional[int]=None, pbar:Optional[PBar]=None,
-                  ordered:bool=False) -> List[Tensor]:
-        "Return predictions and targets on the valid, train, or test set, depending on `ds_type`."
-        self.model.reset()
-        preds = super().get_preds(ds_type=ds_type, with_loss=with_loss, n_batch=n_batch, pbar=pbar)
-        if ordered and hasattr(self.dl(ds_type), 'sampler'):
-            sampler = [i for i in self.dl(ds_type).sampler]
-            reverse_sampler = np.argsort(sampler)
-            preds[0] = preds[0][reverse_sampler,:] if preds[0].dim() > 1 else preds[0][reverse_sampler]
-            preds[1] = preds[1][reverse_sampler,:] if preds[1].dim() > 1 else preds[1][reverse_sampler]
-        return(preds)
-
 
 def accuracy_fwd(input, targs):
     return accuracy(input[...,0], targs[...,0])
 
 
 def accuracy_bwd(input, targs):
-    return accuracy(input[...,1], targs[...,1])
+    return accuracy(input[...,1], targs[...,1])
+
+
+## Replace code in fastai
+import fastai.text.learner
+fastai.text.learner.convert_weights = convert_weights

fastai_contrib/models.py

@@ -22,13 +22,14 @@ class BiLMModel(nn.Module):
             return torch.stack([fwd_o, bwd_o], dim=len(fwd_o.shape))
 
     def forward(self, input):
-        if len(input) == 3: # sl, bs, tracks
+        if len(input.shape) == 3: # sl, bs, tracks
             f = input[..., 0]
             b = input[..., 1]
-        elif len(input) == 2: # sl, bs - support during classification mode
+        elif len(input.shape) == 2: # sl, bs - support during classification mode
             f = input
             b = torch.flip(input, [0])
-
+        else:
+            raise AttributeError(f"Inorrect size of input, {input.shape}")
         fwd_o = self.fwd_lm(f)
         bwd_o = self.bwd_lm(b)
 
@@ -39,6 +40,45 @@ class BiLMModel(nn.Module):
         self.fwd_lm.reset()
         self.bwd_lm.reset()
 
+
+class BiPoolingLinearClassifier(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()
+
+            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 = 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:
@@ -66,6 +106,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), PoolingLinearClassifier(layers, drops))
+    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), BiPoolingLinearClassifier(layers, drops))
     model.reset()
     return model

ulmfit/pretrain_lm.py

@@ -143,7 +143,7 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
     # compared to standard Adam, we set beta_1 to 0.8
     learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99))
 
-    #learn.true_wd = False
+    learn.true_wd = False
     print("true_wd: ", learn.true_wd)
 
     if bidir:

ulmfit/train_clas.py

@@ -7,8 +7,9 @@ import pickle
 
 import torch
 from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner
-from fastai import fit_one_cycle
-from fastai_contrib.learner import bilm_text_classifier_learner, bilm_learner
+from fastai import fit_one_cycle, accuracy
+from fastai_contrib.data import LanguageModelType
+from fastai_contrib.learner import bilm_text_classifier_learner, bilm_learner, accuracy_fwd, accuracy_bwd
 from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists
 from fastai.text.transform import Vocab
 
@@ -68,20 +69,22 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
                         model_dir/f"{pretrained_fname[0]}.pth",
                         model_dir/f"{pretrained_fname[1]}.pkl")
 
-    data_clas, data_lm = get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct)
-
-    if qrnn:
-        emb_sz, nh, nl = 400, 1550, 3
-    else:
-        emb_sz, nh, nl = 400, 1150, 3
-
     if bidir:
+        print("BiLM")
         classifier_learner = bilm_text_classifier_learner
         lm_learner = bilm_learner
     else:
         classifier_learner = text_classifier_learner
         lm_learner = language_model_learner
 
+    lm_type = LanguageModelType.BiLM if bidir else LanguageModelType.FwdLM
+    data_clas, data_lm = get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_type=lm_type)
+
+    if qrnn:
+        emb_sz, nh, nl = 400, 1550, 3
+    else:
+        emb_sz, nh, nl = 400, 1150, 3
+
     lm_enc_finetuned  = f"{lm_name}_{dataset}_{name}_enc"
     if fine_tune and not (model_dir/f"{lm_enc_finetuned}.pth").exists():
         print('Fine-tuning the language model...')
@@ -91,6 +94,11 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
             pretrained_fnames=pretrained_fname,
             path=model_dir.parent, model_dir=model_dir.name,
             drop_mult=0.3)
+        if bidir:
+            learn.metrics = [accuracy_fwd, accuracy_bwd]
+        else:
+            learn.metrics = [accuracy]
+
         learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7))
         learn.unfreeze()
         learn.fit_one_cycle(10, 1e-3, moms=(0.8, 0.7))
@@ -131,7 +139,7 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
     return results
 
 
-def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct):
+def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_type):
     tmp_dir = dataset_dir / 'tmp'
     tmp_dir.mkdir(exist_ok=True)
     vocab_file = tmp_dir / f'vocab_{lang}.pkl'
@@ -170,7 +178,7 @@ def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct):
         for split in [TRN, VAL, TST]:
             ids[split] = ids[split][:int(len(ids[split]) * ds_pct)]
     data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=ids[TRN],
-                                       valid_ids=ids[VAL], bs=bs, bptt=bptt)
+                                       valid_ids=ids[VAL], bs=bs, bptt=bptt, lm_type=lm_type)
     #  TODO TextClasDataBunch allows tst_ids as input, but not tst_lbls?
     data_clas = TextClasDataBunch.from_ids(
         path=tmp_dir, vocab=vocab, train_ids=ids[TRN], valid_ids=ids[VAL],