From 80d4d4da29a8860f9c588db2fd873fbe4af2b17b Mon Sep 17 00:00:00 2001 From: Piotr Czapla Date: Sun, 25 Nov 2018 12:37:45 +0100 Subject: [PATCH] Extract params to an experiment data class You can run this as follows: `python -m ulmfit.pretrain_lm --dir-path 'data/wiki/wikitext-2' --qrnn=True train_lm --num_epochs=1` --- tests/test_end_to_end.py | 16 +-- ulmfit/pretrain_lm.py | 292 ++++++++++++++++++++++----------------- 2 files changed, 170 insertions(+), 138 deletions(-) diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 9b80a6b..a956e40 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -50,17 +50,18 @@ def test_ulmfit_default_end_to_end(): test_data, wt2 = get_test_data() lm_name = 'end-to-end-test-default' cuda_id = 0 - results = ulmfit.pretrain_lm.pretrain_lm( + exp = ulmfit.pretrain_lm.Experiment( dir_path=wt2, lang='en', - cuda_id=cuda_id, qrnn=True, subword=False, max_vocab=1000, bs=2, - num_epochs=1, name=lm_name) - assert results['accuracy'] > 0.02 + + exp.train_lm(num_epochs=1) + + assert exp.results['accuracy'] > 0.02 results = ulmfit.train_clas.new_train_clas( data_dir=test_data, @@ -82,7 +83,7 @@ def test_ulmfit_sentencepiece_end_to_end(): imdb, wt2 = get_test_data() lm_name = 'end-to-end-test-spm' cuda_id = 0 - results = ulmfit.pretrain_lm.pretrain_lm( + exp = ulmfit.pretrain_lm.Experiment( dir_path=wt2, lang='en', cuda_id=cuda_id, @@ -90,11 +91,10 @@ def test_ulmfit_sentencepiece_end_to_end(): subword=True, max_vocab=100, bs=2, - num_epochs=1, name=lm_name, ) - - assert results['accuracy'] > 0.30 + exp.train_lm(num_epochs=1) + assert exp.results['accuracy'] > 0.30 # NOTE: ds_pct is not available for sentencepiece -- tests are on the complete dataset # sentencepiece for finetuning/classification is currently not implemented diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index aaca9e7..4ac57aa 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -25,159 +25,191 @@ import fastai_contrib.data as contrib_data # cupy needs to be installed for QRNN -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): - """ - :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 qrnn: Use a QRNN. Requires installing cupy. - :param subword: Use sub-word tokenization on the cleaned data. - :param max_vocab: The maximum size of the vocabulary. - :param bs: The batch size. - :param bptt: The back-propagation-through-time sequence length. - :param name: The name used for both the model and the vocabulary. - :param model_dir: The path to the directory where the models should be saved - :param bidir: whether the language model is bidirectional - """ - results = {} +# """ +# :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 qrnn: Use a QRNN. Requires installing cupy. +# :param subword: Use sub-word tokenization on the cleaned data. +# :param max_vocab: The maximum size of the vocabulary. +# :param bs: The batch size. +# :param bptt: The back-propagation-through-time sequence length. +# :param name: The name used for both the model and the vocabulary. +# :param model_dir: The path to the directory where the models should be saved +# :param bidir: whether the language model is bidirectional +# """ - if not torch.cuda.is_available(): - print('CUDA not available. Setting device=-1.') - cuda_id = -1 - torch.cuda.set_device(cuda_id) - dir_path = Path(dir_path) - assert dir_path.exists() - model_dir = dir_path / 'models' # removed from params, as it is absolute models location in train_clas and here it is relative - model_dir.mkdir(exist_ok=True) - print('Batch size:', bs) - print('Max vocab:', max_vocab) - model_name = 'qrnn' if qrnn else 'lstm' - if qrnn: - print('Using QRNNs...') +@dataclass +class Experiment: + dir_path: str + bidir: bool =False + bptt: int = 70 + bs: int = 70 + lang: str = 'en' + max_vocab: int = 60000 + name: str = 'wt-103' + subword: bool = False + ds_pct: float = 1.0 + qrnn: bool = True + cuda_id:int = 0 + def __post_init__(self): + self.results = {} - trn_path = dir_path / f'{lang}.wiki.train.tokens' - val_path = dir_path / f'{lang}.wiki.valid.tokens' - tst_path = dir_path / f'{lang}.wiki.test.tokens' - for path_ in [trn_path, val_path, tst_path]: - assert path_.exists(), f'Error: {path_} does not exist.' + if not torch.cuda.is_available(): + print('CUDA not available. Setting device=-1.') + cuda_id = -1 + torch.cuda.set_device(self.cuda_id) - if subword: - # apply sentencepiece tokenization - trn_path = dir_path / f'{lang}.wiki.train.tokens' - val_path = dir_path / f'{lang}.wiki.valid.tokens' + self.dir_path = Path(self.dir_path) + assert self.dir_path.exists() + self.model_dir = self.dir_path / 'models' # removed from params, as it is absolute models location in train_clas and here it is relative + self.model_dir.mkdir(exist_ok=True) + print('Batch size:', self.bs) + print('Max vocab:', self.max_vocab) - read_file(trn_path, 'train') - read_file(val_path, 'valid') + if self.qrnn: + print('Using QRNNs...') - sp = get_sentencepiece(dir_path, trn_path, name, vocab_size=max_vocab) + @property + def model_name(self): + return 'qrnn' if self.qrnn else 'lstm' - lm_type = contrib_data.LanguageModelType.BiLM if bidir else contrib_data.LanguageModelType.FwdLM + def train_lm(self, num_epochs=10): + data_lm = self.load_data() + exe = Executor(self.create_lm_learner(data_lm), exp=self) + if num_epochs > 0: + exe.learn.fit_one_cycle(num_epochs, 5e-3, (0.8, 0.7), wd=1e-7) + exe.validate() + exe.save() + return exe - data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp, bs=bs, bptt=bptt, lm_type=lm_type) - itos = data_lm.train_ds.vocab.itos - stoi = data_lm.train_ds.vocab.stoi - else: - # read the already whitespace separated data without any preprocessing - trn_tok = read_whitespace_file(trn_path) - val_tok = read_whitespace_file(val_path) - if ds_pct < 1.0: - trn_tok = trn_tok[:max(20, int(len(trn_tok) * ds_pct))] - val_tok = val_tok[:max(20, int(len(val_tok) * ds_pct))] - print(f"Limiting data sets to {ds_pct*100}%, trn {len(trn_tok)}, val: {len(val_tok)}") - - itos_fname = model_dir / f'itos_{name}.pkl' - if not itos_fname.exists(): - # create the vocabulary - cnt = Counter(word for sent in trn_tok for word in sent) - itos = [o for o,c in cnt.most_common(n=max_vocab)] - itos.insert(1, PAD) #  set pad id to 1 to conform to fast.ai standard - assert UNK in itos, f'Unknown words are expected to have been replaced with {UNK} in the data.' - - # save vocabulary - print(f"Saving vocabulary as {itos_fname}") - results['itos_fname'] = itos_fname - with open(itos_fname, 'wb') as f: - pickle.dump(itos, f) + def create_lm_learner(self, data_lm): + # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103) + # for training on smaller datasets, more dropout is necessary + if self.qrnn: + emb_sz, nh, nl = 400, 1550, 3 + # dps = np.array([0.0, 0.0, 0.0, 0.0, 0.0]) + dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) + drop_mult = 0.1 else: - print("Loading itos:", itos_fname) - itos = np.load(itos_fname) - vocab = Vocab(itos) - stoi = vocab.stoi + emb_sz, nh, nl = 400, 1150, 3 + # emb_sz, nh, nl = 400, 1150, 3 + dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) + drop_mult = 0.1 + fastai.text.learner.default_dropout['language'] = dps + lm_learner = bilm_learner if self.bidir else language_model_learner + learn = lm_learner(data_lm, bptt=self.bptt, emb_sz=emb_sz, nh=nh, nl=nl, pad_token=1, + drop_mult=drop_mult, tie_weights=True, model_dir=self.model_dir.name, + bias=True, qrnn=self.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 + print("true_wd: ", learn.true_wd) + if self.bidir: + learn.metrics = [accuracy_fwd, accuracy_bwd] + else: + learn.metrics = [accuracy] + return learn - trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok]) - val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_tok]) - lm_type = contrib_data.LanguageModelType.BiLM if bidir else contrib_data.LanguageModelType.FwdLM + def load_data(self): + trn_path = self.dir_path / f'{self.lang}.wiki.train.tokens' + val_path = self.dir_path / f'{self.lang}.wiki.valid.tokens' + tst_path = self.dir_path / f'{self.lang}.wiki.test.tokens' + for path_ in [trn_path, val_path, tst_path]: + assert path_.exists(), f'Error: {path_} does not exist.' + if self.subword: + # apply sentencepiece tokenization + trn_path = self.dir_path / f'{self.lang}.wiki.train.tokens' + val_path = self.dir_path / f'{self.lang}.wiki.valid.tokens' - # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos)) - data_lm = TextLMDataBunch.from_ids(path=dir_path, vocab=vocab, train_ids=trn_ids, - valid_ids=val_ids, bs=bs, bptt=bptt, - lm_type=lm_type - ) + read_file(trn_path, 'train') + read_file(val_path, 'valid') - print('Size of vocabulary:', len(itos)) - print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)])) + sp = get_sentencepiece(self.dir_path, trn_path, self.name, vocab_size=self.max_vocab) - # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103) - # for training on smaller datasets, more dropout is necessary - if qrnn: - emb_sz, nh, nl = 400, 1550, 3 - #dps = np.array([0.0, 0.0, 0.0, 0.0, 0.0]) - dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) - drop_mult = 0.1 - else: - emb_sz, nh, nl = 400, 1150, 3 - # emb_sz, nh, nl = 400, 1150, 3 - dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) - drop_mult = 0.1 + lm_type = contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM - fastai.text.learner.default_dropout['language'] = dps + data_lm = TextLMDataBunch.from_csv(self.dir_path, 'train.csv', **sp, bs=self.bs, bptt=self.bptt, lm_type=lm_type) + itos = data_lm.train_ds.vocab.itos + stoi = data_lm.train_ds.vocab.stoi + else: + # read the already whitespace separated data without any preprocessing + trn_tok = read_whitespace_file(trn_path) + val_tok = read_whitespace_file(val_path) + if self.ds_pct < 1.0: + trn_tok = trn_tok[:max(20, int(len(trn_tok) * self.ds_pct))] + val_tok = val_tok[:max(20, int(len(val_tok) * self.ds_pct))] + print(f"Limiting data sets to {self.ds_pct * 100}%, trn {len(trn_tok)}, val: {len(val_tok)}") - lm_learner = bilm_learner if bidir else language_model_learner - learn = lm_learner(data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, pad_token=1, - drop_mult=drop_mult, tie_weights=True, model_dir=model_dir.name, - 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)) + itos_fname = self.model_dir / f'itos_{self.name}.pkl' + if not itos_fname.exists(): + # create the vocabulary + cnt = Counter(word for sent in trn_tok for word in sent) + itos = [o for o, c in cnt.most_common(n=self.max_vocab)] + itos.insert(1, PAD) #   set pad id to 1 to conform to fast.ai standard + assert UNK in itos, f'Unknown words are expected to have been replaced with {UNK} in the data.' - learn.true_wd = False - print("true_wd: ", learn.true_wd) + # save vocabulary + print(f"Saving vocabulary as {itos_fname}") + self.results['itos_fname'] = itos_fname + with open(itos_fname, 'wb') as f: + pickle.dump(itos, f) + else: + print("Loading itos:", itos_fname) + itos = np.load(itos_fname) + vocab = Vocab(itos) + stoi = vocab.stoi - if bidir: - learn.metrics = [accuracy_fwd, accuracy_bwd] - else: - learn.metrics = [accuracy] + trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok]) + val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_tok]) - try: - learn.load(f'{model_name}_{name}') - print("Weights loaded") - except FileNotFoundError: - print("Starting from random weights") - pass + lm_type = contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM - learn.fit_one_cycle(num_epochs, 5e-3, (0.8, 0.7), wd=1e-7) + # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos)) + data_lm = TextLMDataBunch.from_ids(path=self.dir_path, vocab=vocab, train_ids=trn_ids, + valid_ids=val_ids, bs=self.bs, bptt=self.bptt, + lm_type=lm_type) + itos, stoi, trn_path = data_lm.vocab.itos, data_lm.vocab.stoi, data_lm.path + print('Size of vocabulary:', len(itos)) + print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)])) + return data_lm - if not subword 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(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}') +class Executor: + def __init__(self, learn, exp): + self.exp = exp + self.learn = learn + try: + self.learn.load(f'{self.exp.model_name}_{self.exp.name}') + print("Weights loaded") + except FileNotFoundError: + print("Starting from random weights") + pass - 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) + def validate(self): + if not self.exp.subword and self.exp.max_vocab is None: + raise NotImplementedError("figure out how to validate and save results") + # 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(learn.model, tst_ids, self.exp.bptt) + print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item()) + + def save(self): + print(f"Saving models at {self.learn.path / self.learn.model_dir}") + self.learn.save(f'{self.exp.model_name}_{self.exp.name}') + + opt_state_path = self.learn.path / self.learn.model_dir / f'{self.exp.model_name}_{self.exp.name}_state.pth' + print(f"Saving optimiser state at {opt_state_path}") + torch.save(self.learn.opt.opt.state_dict(), opt_state_path) + + self.exp.results['accuracy'] = self.learn.validate()[1] #TODO rewrite - results['accuracy'] = learn.validate()[1] - return results if __name__ == '__main__': - fire.Fire(pretrain_lm) + fire.Fire(Experiment)