From 70680c21ce808de6185609b8c8be353b87ab4c38 Mon Sep 17 00:00:00 2001 From: Nirant K Date: Wed, 14 Nov 2018 13:54:31 +0000 Subject: [PATCH] Add better logs, add multi-language flag to pretrain - Added lang variable to pretrain_lm - Tried the QRNN model with Hindi - Added better logging and consistent model saving paths - Fixed minor bug in validate(model, ...) to to validate(learn.model, ...) --- ulmfit/pretrain_lm.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index ffa4503..a4b2f4e 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -25,10 +25,11 @@ from collections import Counter # cupy needs to be installed for QRNN -def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, +def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab=60000, bs=70, bptt=70, name='wt-103', model_dir='models', num_epochs=10): """ :param dir_path: The path to the directory of the file. + :param lang: the language unicode :param cuda_id: The id of the GPU. Uses GPU 0 by default or no GPU when run on CPU. :param qrrn: Use a QRNN. Requires installing cupy. @@ -54,9 +55,9 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, if qrnn: print('Using QRNNs...') - trn_path = dir_path / 'wiki.train.tokens' - val_path = dir_path / 'wiki.valid.tokens' - tst_path = dir_path / 'wiki.test.tokens' + trn_path = dir_path / f'{lang}.wiki.train.tokens.unk' + val_path = dir_path / f'{lang}.wiki.valid.tokens.unk' + tst_path = dir_path / f'{lang}.wiki.test.tokens.unk' for path_ in [trn_path, val_path, tst_path]: assert path_.exists(), f'Error: {path_} does not exist.' @@ -107,15 +108,15 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, fastai.text.learner.default_dropout['language'] = dps learn = language_model_learner(data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, pad_token=1, - drop_mult=drop_mult, tie_weights=True, + drop_mult=drop_mult, tie_weights=True, model_dir=model_dir, bias=True, qrnn=True, 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 # save vocabulary - print('Saving vocabulary...') - with open(model_dir / f'itos_{name}.pkl', 'wb') as f: + print(f"Saving vocabulary as {dir_path / model_dir}") + with open(dir_path / model_dir / f'itos_{name}.pkl', 'wb') as f: pickle.dump(itos, f) fit_one_cycle(learn, num_epochs, 5e-3, (0.8, 0.7), wd=1e-7) @@ -123,13 +124,18 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, if clean and max_vocab is None: # only if we use the unpreprocessed version and the full vocabulary # are the perplexity results comparable to previous work + print(f"Validating model performance with test tokens from: {trn_path}") tst_tok = read_whitespace_file(trn_path) tst_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in tst_tok]) - logloss, perplexity = validate(model, tst_ids, bptt) + logloss, perplexity = validate(learn.model, tst_ids, bptt) print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item()) + print(f"Saving models at {learn.path / learn.model_dir}") learn.save(f'{model_name}_{name}') - torch.save(learn.opt.opt.state_dict(), learn.path / learn.model_dir / f'{model_name}3_{name}_state.pth') + + opt_state_path = learn.path / learn.model_dir / f'{model_name}3_{name}_state.pth' + print(f"Saving optimiser state at {opt_state_path}") + torch.save(learn.opt.opt.state_dict(), opt_state_path) if __name__ == '__main__':