diff --git a/fastai_contrib/learner.py b/fastai_contrib/learner.py index d787867..3b9bafd 100644 --- a/fastai_contrib/learner.py +++ b/fastai_contrib/learner.py @@ -1,7 +1,8 @@ +from fastai import GradientClipping, accuracy from fastai.callbacks import * from fastai.basic_data import * from fastai.datasets import untar_data -from fastai_contrib.models import get_bilm +from fastai_contrib.models import get_bilm, get_rnn_classifier, get_birnn_classifier from fastai.text.learner import * @@ -26,7 +27,88 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in return learn +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, + drop_mult: float = 1., qrnn: bool = False, **kwargs) -> 'TextClassifierLearner': + "Create a RNN classifier." + dps = default_dropout['classifier'] * drop_mult + if lin_ftrs is None: lin_ftrs = [50] + 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] + 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], + qrnn=qrnn) + learn = RNNLearner(data, model, bptt, split_func=birnn_classifier_split, **kwargs) + return learn + + def bilm_split(model:nn.Module) -> List[nn.Module]: "Split a RNN `model` in groups for differential learning rates." return [f+b for f,b in zip(lm_split(model.fwd_lm),lm_split(model.bwd_lm))] + +def birnn_classifier_split(model:nn.Module) -> List[nn.Module]: + "Split a RNN `model` in groups for differential learning rates." + f_rnn,b_rnn = model[0].fwd_lm,model[0].bwd_lm + groups = [[f_rnn.encoder, f_rnn.encoder_dp,b_rnn.encoder, b_rnn.encoder_dp]] + groups += [a for a in zip(f_rnn.rnns, f_rnn.hidden_dps, b_rnn.rnns, b_rnn.hidden_dps, )] + 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]) \ No newline at end of file diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py index df323af..2b420e7 100644 --- a/fastai_contrib/models.py +++ b/fastai_contrib/models.py @@ -9,17 +9,30 @@ class BiLMModel(nn.Module): self.fwd_lm = fwd_lm self.bwd_lm = bwd_lm + def __getitem__(self, idx): + return BiLMModel(self.fwd_lm[idx], self.bwd_lm[idx]) + + def __len__(self): + return len(self.fwd_lm) + + def stack(self, fwd_o, bwd_o): + if is_listy(fwd_o): + return [self.stack(f, b) for f,b in zip(fwd_o,bwd_o)] + else: + return torch.stack([fwd_o, bwd_o], dim=len(fwd_o.shape)) + def forward(self, input): - sl, bs, tracks = input.size() + if len(input) == 3: # sl, bs, tracks + f = input[..., 0] + b = input[..., 1] + elif len(input) == 2: # sl, bs - support during classification mode + f = input + b = torch.flip(input, [0]) - decoded = [] - raw_outputs = [] - outputs = [] + fwd_o = self.fwd_lm(f) + bwd_o = self.bwd_lm(b) - fwd_o = self.fwd_lm(input[..., 0]) - bwd_o = self.bwd_lm(input[..., 1]) - - return torch.stack([fwd_o[0], bwd_o[0]], dim=2), (fwd_o[1]+bwd_o[1]), (fwd_o[2] + bwd_o[2]) + return self.stack(fwd_o, bwd_o) def reset(self): "Reset the hidden states of underlaying lms." @@ -42,4 +55,17 @@ def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, t return BiLMModel( fwd_lm=SequentialRNN(fwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)), - bwd_lm=SequentialRNN(bwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias))) \ No newline at end of file + bwd_lm=SequentialRNN(bwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias))) + +def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, + pad_token:int, layers:Collection[int], drops:Collection[float], bidir:bool=False, qrnn:bool=False, + hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5)->nn.Module: + "Create a RNN classifier model." + fwd_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) + 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.reset() + return model \ No newline at end of file diff --git a/results/bilm_qrnn-wt-103-unk-wd1 b/results/bilm_qrnn-wt-103-unk-wd1 new file mode 100644 index 0000000..dc5c8ce --- /dev/null +++ b/results/bilm_qrnn-wt-103-unk-wd1 @@ -0,0 +1,20 @@ +(fastaiv1) n-waves@GV100:~/workspace/ulmfit-multilingual$ python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --cuda-id=0 - +-name=bilm-wt-103-unk --num_epochs=10 --bs=64 --bidir=True +Batch size: 64 +Max vocab: 60000 +Using QRNNs... +Loading itos: data/wiki/wikitext-103-unk/models/itos_bilm-wt-103-unk.pkl +Size of vocabulary: 60001 +First 10 words in vocab: the, , ,, ., of, and, to, in, , a +Starting from random weights +epoch train_loss valid_loss accuracy_fwd accuracy_bwd +1 4.184118 4.113817 0.308417 0.318290 +2 4.106002 4.039978 0.312220 0.320315 +3 4.160282 4.086650 0.306327 0.314152 +4 4.136789 4.049762 0.309202 0.317373 +5 4.086520 4.001534 0.313158 0.321543 +6 4.036415 3.960494 0.318333 0.325623 +7 3.988382 3.913568 0.324409 0.330398 +8 3.937409 3.874283 0.328386 0.334786 +9 3.912215 3.856325 0.330593 0.337102 +10 3.876535 3.848783 0.331202 0.337649 \ No newline at end of file diff --git a/results/bilm_qrnn-wt-103-wd0 b/results/bilm_qrnn-wt-103-wd0 new file mode 100644 index 0000000..291173e --- /dev/null +++ b/results/bilm_qrnn-wt-103-wd0 @@ -0,0 +1,20 @@ +(fastaiv1) n-waves@GV100:~/workspace/ulmfit-multilingual$ python -m ulmfit.pretrain_lm +ki/wikitext-103 --qrnn=True --cuda-id=0 --name=bilm-wt-103 --num_epochs=10 --bs=64 --bidir=True +Batch size: 64 +Max vocab: 60000 +Using QRNNs... +Loading itos: data/wiki/wikitext-103/models/itos_bilm-wt-103.pkl +Size of vocabulary: 60001 +First 10 words in vocab: the, , ,, ., of, and, to, in, , a +Starting from random weights +epoch train_loss valid_loss accuracy_fwd accuracy_bwd +1 4.141750 4.114931 0.307700 0.320260 +2 4.078711 4.061186 0.310221 0.320550 +3 4.108722 4.081286 0.308383 0.316431 +4 4.054437 4.040497 0.311229 0.319658 +5 4.014880 3.966010 0.318588 0.327607 +6 3.926770 3.897305 0.326874 0.333171 +7 3.855364 3.816452 0.335632 0.342045 +8 3.761600 3.745798 0.343413 0.350675 + +10 3.614097 3.687730 0.351651 0.358280 \ No newline at end of file diff --git a/results/lm_lstm-wt-103-wd0 b/results/lm_lstm-wt-103-wd0 new file mode 100644 index 0000000..41c178f --- /dev/null +++ b/results/lm_lstm-wt-103-wd0 @@ -0,0 +1,22 @@ +(fastaiv1) pczapla@galatea ~/w/ulmfit-multilingual ❯❯❯ time python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=False --cuda-id=1 --name=wt-103 --num_epochs=10 --bs=32 +Batch size: 32 +Max vocab: 60000 +Loading itos: data/wiki/wikitext-103/models/itos_wt-103.pkl +Size of vocabulary: 60001 +First 10 words in vocab: the, , ,, ., of, and, to, in, , a +Starting from random weights +epoch train_loss valid_loss accuracy +1 4.317837 4.245714 0.309034 +2 4.347099 4.251623 0.306932 +3 4.360198 4.302850 0.302014 +4 4.329998 4.264225 0.306828 +5 4.301852 4.209807 0.311763 +6 4.218213 4.131647 0.319867 +7 4.167365 4.047796 0.327259 +8 4.095695 3.980298 0.336255 +9 4.032371 3.919622 0.343671 +10 3.983160 3.906227 0.345863 +Saving models at data/wiki/wikitext-103/models +Saving optimiser state at data/wiki/wikitext-103/models/lstm3_wt-103_state.pth +accuracy: tensor(0.3461) +python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=False --cuda-id=1 44147.64s user 16746.69s system 99% cpu 16:58:47.57 total \ No newline at end of file diff --git a/results/lm_qrnn-wt-103-unk-wd0 b/results/lm_qrnn-wt-103-unk-wd0 new file mode 100644 index 0000000..d009387 --- /dev/null +++ b/results/lm_qrnn-wt-103-unk-wd0 @@ -0,0 +1,23 @@ +time python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --cuda-id=0 --name=wt-103 --num_epochs=10 --bs=32 +Batch size: 32 +Max vocab: 60000 +Using QRNNs... +Loading itos: data/wiki/wikitext-103-unk/models/itos_wt-103.pkl +Size of vocabulary: 60001 +First 10 words in vocab: the, , ,, ., of, and, to, in, , a +Starting from random weights +epoch train_loss valid_loss accuracy +1 4.393613 4.282537 0.305135 +2 4.394575 4.299973 0.300028 +3 4.444056 4.336682 0.295496 +4 4.414292 4.332908 0.297408 +5 4.406248 4.269964 0.302962 +6 4.327769 4.209948 0.309025 +7 4.244819 4.140769 0.315675 +8 4.210000 4.075574 0.323544 +9 4.139524 4.034881 0.328550 +10 4.150556 4.021361 0.331256 +Saving models at data/wiki/wikitext-103-unk/models +Saving optimiser state at data/wiki/wikitext-103-unk/models/qrnn3_wt-103_state.pth +accuracy: tensor(0.3312) +python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --bs=3 32119.72s user 12439.18s system 99% cpu 12:25:03.65 total \ No newline at end of file diff --git a/results/lm_qrnn-wt-103-wd0 b/results/lm_qrnn-wt-103-wd0 new file mode 100644 index 0000000..33d243b --- /dev/null +++ b/results/lm_qrnn-wt-103-wd0 @@ -0,0 +1,20 @@ +time python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=True --cuda-id=1 --name=wt-103 --num_epochs=10 --bs=32 ✘ 1 +Batch size: 32 +Max vocab: 60000 +Using QRNNs... +Size of vocabulary: 60001 +First 10 words in vocab: the, , ,, ., of, and, to, in, , a +Saving vocabulary as data/wiki/wikitext-103/models +epoch train_loss valid_loss accuracy +1 4.328219 4.268609 0.306414 +2 4.351323 4.282290 0.300560 +3 4.376134 4.345460 0.294145 +4 4.394314 4.319661 0.298142 +5 4.337362 4.265630 0.303437 +6 4.295763 4.199119 0.309354 +7 4.165526 4.121037 0.318062 +8 4.151789 4.060387 0.325899 +9 4.066542 4.014568 0.332732 +10 4.013142 3.997691 0.335299 +Saving models at data/wiki/wikitext-103/models +Saving optimiser state at data/wiki/wikitext-103/models/qrnn3_wt-103_state.pth \ No newline at end of file diff --git a/tests/test_text_train.py b/tests/test_text_train.py index 5116cfa..0103939 100644 --- a/tests/test_text_train.py +++ b/tests/test_text_train.py @@ -4,9 +4,11 @@ from fastai.text import * pytestmark = pytest.mark.integration +print(sys.path) import fastai_contrib.data as contrib_data -from fastai_contrib.learner import bilm_learner +from fastai_contrib.learner import bilm_learner, accuracy_fwd, bilm_text_classifier_learner + def read_file(fname): texts = [] @@ -38,11 +40,44 @@ def learn(): learn.fit_one_cycle(4, 5e-3) return learn +def text_df(n_labels): + data = [] + texts = ["fast ai is a cool project", "hello world"] + for ind, text in enumerate(texts): + sample = {} + for label in range(n_labels): sample[label] = ind%2 + sample["text"] = text + data.append(sample) + df = pd.DataFrame(data) + return df + ###################### NEW CODE def test_val_loss(learn): assert learn.validate()[1] > 0.5 + +def test_bilm_classifier_loads_encoder(): + n_labels=2 + path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'tmp') + os.makedirs(path) + try: + df = text_df(n_labels=1) + lmdf = df#[["text"]] + print(lmdf.head()) + lmdata = TextLMDataBunch.from_df(path, lmdf, lmdf, tokenizer=Tokenizer(BaseTokenizer), + lm_type=contrib_data.LanguageModelType.BiLM) + learn = bilm_learner(lmdata, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False) + learn.save_encoder("enc") + data = TextClasDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=list(range(n_labels)), text_cols=["text"]) + classifier = bilm_text_classifier_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False) + print(last_layer(classifier.model), ) + classifier.load_encoder("enc") + classifier.fit(1) + finally: + shutil.rmtree(path) + + def test_bilm_lstm_can_be_trained(): manual_seed() path, df_trn, df_val = prep_human_numbers() @@ -50,9 +85,10 @@ def test_bilm_lstm_can_be_trained(): lm_type = contrib_data.LanguageModelType.BiLM) learn = bilm_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False) - learn.metrics = [] - learn.fit_one_cycle(4, 5e-3) - assert learn.validate()[0] < 2 #TODO Change to accuracy once it is fixed + learn.metrics = [accuracy_fwd] + learn.fit_one_cycle(2, 5e-3) + assert learn.validate()[1] > 0.3 + def test_bwdlm_lstm_can_be_trained(): manual_seed() @@ -61,5 +97,5 @@ def test_bwdlm_lstm_can_be_trained(): lm_type = contrib_data.LanguageModelType.BwdLM) learn = language_model_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False) - learn.fit_one_cycle(4, 5e-3) - assert learn.validate()[1] > 0.5 + learn.fit_one_cycle(2, 5e-3) + assert learn.validate()[1] > 0.3 diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index cc1ecdc..856fc82 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -6,14 +6,13 @@ to be split. """ import fastai import fire -import numpy as np from fastai import * from fastai.text import * import torch -from fastai_contrib.utils import read_file, read_whitespace_file,\ - DataStump, validate, PAD, UNK, get_sentencepiece -from fastai_contrib.learner import bilm_learner +from fastai_contrib.utils import read_file, read_whitespace_file, \ + validate, PAD, UNK, get_sentencepiece +from fastai_contrib.learner import bilm_learner, accuracy_fwd, accuracy_bwd import pickle from pathlib import Path @@ -25,11 +24,6 @@ import fastai_contrib.data as contrib_data # conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92] # cupy needs to be installed for QRNN -def accuracy_fwd(input, targs): - return accuracy(input[...,0], targs[...,0]) -def accuracy_bwd(input, targs): - return accuracy(input[...,1], targs[...,1]) - def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vocab=60000, bs=70, bptt=70, name='wt-103', num_epochs=10, bidir=False, ds_pct=1.0): @@ -148,7 +142,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo bias=True, qrnn=qrnn, clip=0.12) # 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: learn.metrics = [accuracy_fwd, accuracy_bwd] diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py index 27f014e..747c421 100644 --- a/ulmfit/train_clas.py +++ b/ulmfit/train_clas.py @@ -8,6 +8,7 @@ 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_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists from fastai.text.transform import Vocab @@ -19,7 +20,7 @@ from pathlib import Path def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_dir='models', qrnn=False, fine_tune=True, max_vocab=60000, bs=20, bptt=70, name='imdb-clas', - dataset='imdb', ds_pct=1.0): + dataset='imdb', bidir=False, ds_pct=1.0, train=True): """ :param data_dir: The path to the `data` directory :param lang: the language unicode @@ -67,11 +68,73 @@ 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: + classifier_learner = bilm_text_classifier_learner + lm_learner = bilm_learner + else: + classifier_learner = text_classifier_learner + lm_learner = language_model_learner + + 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...') + learn = lm_learner( + data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, qrnn=qrnn, + pad_token=PAD_TOKEN_ID, + pretrained_fnames=pretrained_fname, + path=model_dir.parent, model_dir=model_dir.name, + drop_mult=0.3) + 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)) + + # save encoder + learn.save_encoder(lm_enc_finetuned) + + + learn = classifier_learner(data_clas, bptt=bptt, pad_token=PAD_TOKEN_ID, + path=model_dir.parent, model_dir=model_dir.name, + qrnn=qrnn, emb_sz=emb_sz, nh=nh, nl=nl, drop_mult=0.5) + + try: + learn.load(f'{model_name}_{name}') + print("Loading classifier") + except FileNotFoundError: + learn.load_encoder(lm_enc_finetuned) + print("loading encoder") + train = True + + if train: + print("Starting classifier training") + learn.fit_one_cycle(1, 2e-2, moms=(0.8, 0.7), wd=1e-7) + + learn.freeze_to(-2) + learn.fit_one_cycle(1, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7)) + + learn.freeze_to(-3) + learn.fit_one_cycle(1, slice(5e-3 / (2.6 ** 4), 5e-3), moms=(0.8, 0.7)) + + learn.unfreeze() + learn.fit_one_cycle(2, slice(1e-3 / (2.6 ** 4), 1e-3), moms=(0.8, 0.7)) + + print(f"Saving models at {learn.path / learn.model_dir}") + learn.save(f'{model_name}_{name}') + + results['accuracy'] = learn.validate()[1] + return results + + +def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct): tmp_dir = dataset_dir / 'tmp' tmp_dir.mkdir(exist_ok=True) vocab_file = tmp_dir / f'vocab_{lang}.pkl' - if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists(): print('Reading the data...') toks, lbls = read_clas_data(dataset_dir, dataset, lang) @@ -100,66 +163,20 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_ lbls[split] = np.load(tmp_dir / f'{split}_{lang}_lbl.npy') with open(vocab_file, 'rb') as f: vocab = pickle.load(f) - print(f'Train size: {len(ids[TRN])}. Valid size: {len(ids[VAL])}. ' f'Test size: {len(ids[TST])}.') - if ds_pct < 1.0: print(f"Making the dataset smaller {ds_pct}") for split in [TRN, VAL, TST]: - ids[split] = ids[split][:int(len(ids[split])*ds_pct)] - - + 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) - - # TODO TextClasDataBunch allows tst_ids as input, but not tst_lbls? + #  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], train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs) + return data_clas, data_lm - if qrnn: - emb_sz, nh, nl = 400, 1550, 3 - else: - emb_sz, nh, nl = 400, 1150, 3 - learn = language_model_learner( - data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, qrnn=qrnn, - pad_token=PAD_TOKEN_ID, - pretrained_fnames=pretrained_fname, - path=model_dir.parent, model_dir=model_dir.name, - drop_mult=0.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...') - 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)) - - # save encoder - learn.save_encoder(lm_enc_finetuned) - - print("Starting classifier training") - learn = text_classifier_learner(data_clas, bptt=bptt, pad_token=PAD_TOKEN_ID, - path=model_dir.parent, model_dir=model_dir.name, - qrnn=qrnn, emb_sz=emb_sz, nh=nh, nl=nl, drop_mult=0.5) - - learn.load_encoder(lm_enc_finetuned) - - learn.fit_one_cycle(1, 2e-2, moms=(0.8, 0.7), wd=1e-7) - - learn.freeze_to(-2) - learn.fit_one_cycle(1, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7)) - - learn.freeze_to(-3) - learn.fit_one_cycle(1, slice(5e-3 / (2.6 ** 4), 5e-3), moms=(0.8, 0.7)) - - learn.unfreeze() - learn.fit_one_cycle(2, slice(1e-3 / (2.6 ** 4), 1e-3), moms=(0.8, 0.7)) - results['accuracy'] = learn.validate()[1] - print(f"Saving models at {learn.path / learn.model_dir}") - learn.save(f'{model_name}_{name}') - return results if __name__ == '__main__': fire.Fire(new_train_clas)