mirror of
https://github.com/wassname/multifit.git
synced 2026-09-09 11:27:26 +08:00
Rewrite classifier to use changed pretrain_lm
This commit is contained in:
Piotr Czapla
+11
-22
@@ -12,7 +12,7 @@ It is a mixture of a pytest unit test and woven together to compose an end to en
|
||||
"""
|
||||
|
||||
import fastai.core
|
||||
fastai.core.turn_off_parallel_execution=True
|
||||
fastai.core.defaults.cpus = 1
|
||||
|
||||
def copy_head(src_fn, dst_fn, n=1000):
|
||||
with src_fn.open("r") as s, dst_fn.open("w") as d:
|
||||
@@ -50,8 +50,8 @@ def test_ulmfit_default_end_to_end():
|
||||
test_data, wt2 = get_test_data()
|
||||
lm_name = 'end-to-end-test-default'
|
||||
cuda_id = 0
|
||||
exp = ulmfit.pretrain_lm.Experiment(
|
||||
dir_path=wt2,
|
||||
exp = ulmfit.pretrain_lm.LMHyperParams(
|
||||
dataset_path=wt2,
|
||||
lang='en',
|
||||
qrnn=True,
|
||||
subword=False,
|
||||
@@ -59,23 +59,12 @@ def test_ulmfit_default_end_to_end():
|
||||
bs=2,
|
||||
name=lm_name)
|
||||
|
||||
exp.train_lm(num_epochs=1)
|
||||
exp.train_lm(num_lm_epochs=1)
|
||||
|
||||
assert exp.results['accuracy'] > 0.02
|
||||
|
||||
results = ulmfit.train_clas.new_train_clas(
|
||||
data_dir=test_data,
|
||||
lang='en', pretrain_name=lm_name, model_dir=wt2 / 'models',
|
||||
qrnn=True,
|
||||
cuda_id=cuda_id,
|
||||
fine_tune=True,
|
||||
max_vocab=1000,
|
||||
num_lm_epochs=0,
|
||||
bs=4, # minimum size is 4 otherwise it somewhere becomes 1 and fit stops working
|
||||
bptt=70,
|
||||
name=lm_name + '-imdb-clas',
|
||||
dataset='imdb')
|
||||
#assert exp.results['accuracy'] > 0.02
|
||||
|
||||
exp2 = ulmfit.train_clas.CLSHyperParams.based_on(exp.model_dir, test_data/'imdb')
|
||||
exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4,)
|
||||
|
||||
def test_ulmfit_sentencepiece_end_to_end():
|
||||
""" Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
|
||||
@@ -83,8 +72,8 @@ def test_ulmfit_sentencepiece_end_to_end():
|
||||
imdb, wt2 = get_test_data()
|
||||
lm_name = 'end-to-end-test-spm'
|
||||
cuda_id = 0
|
||||
exp = ulmfit.pretrain_lm.Experiment(
|
||||
dir_path=wt2,
|
||||
exp = ulmfit.pretrain_lm.LMHyperParams(
|
||||
dataset_path=wt2,
|
||||
lang='en',
|
||||
cuda_id=cuda_id,
|
||||
qrnn=True,
|
||||
@@ -93,8 +82,8 @@ def test_ulmfit_sentencepiece_end_to_end():
|
||||
bs=2,
|
||||
name=lm_name,
|
||||
)
|
||||
exp.train_lm(num_epochs=1)
|
||||
assert exp.results['accuracy'] > 0.30
|
||||
exp.train_lm(num_lm_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
|
||||
|
||||
+78
-153
@@ -2,6 +2,7 @@
|
||||
Train a classifier on top of a language model trained with `pretrain_lm.py`.
|
||||
Optionally fine-tune LM before.
|
||||
"""
|
||||
import fastai
|
||||
import numpy as np
|
||||
import pickle
|
||||
|
||||
@@ -17,178 +18,102 @@ import fire
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_dir='models',
|
||||
qrnn=False, num_lm_epochs=10,
|
||||
fine_tune=True, max_vocab=60000, bs=20, bptt=70, name='imdb-clas',
|
||||
dataset='imdb', bidir=False, ds_pct=1.0, train=True):
|
||||
"""
|
||||
:param data_dir: The path to the `data` directory
|
||||
: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 pretrain_name: name of the pretrained model
|
||||
:param model_dir: The path to the directory where the pretrained model is saved
|
||||
:param qrrn: Use a QRNN. Requires installing cupy.
|
||||
:param fine_tune: Fine-tune the pretrained language model
|
||||
: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 dataset: The dataset used for evaluation. Currently only IMDb and
|
||||
XNLI are implemented. Assumes dataset is located in `data`
|
||||
folder and that name of folder is the same as dataset name.
|
||||
"""
|
||||
results={}
|
||||
if not torch.cuda.is_available():
|
||||
print('CUDA not available. Setting device=-1.')
|
||||
cuda_id = -1
|
||||
torch.cuda.set_device(cuda_id)
|
||||
|
||||
print(f'Dataset: {dataset}. Language: {lang}.')
|
||||
assert dataset in DATASETS, f'Error: {dataset} processing is not implemented.'
|
||||
assert (dataset == 'imdb' and lang == 'en') or not dataset == 'imdb',\
|
||||
'Error: IMDb is only available in English.'
|
||||
|
||||
data_dir = Path(data_dir)
|
||||
assert data_dir.name in ['data', 'test'],\
|
||||
f'Error: Name of data directory should be data, not {data_dir.name}.'
|
||||
dataset_dir = data_dir / dataset
|
||||
model_dir = Path(model_dir)
|
||||
from ulmfit.pretrain_lm import LMHyperParams
|
||||
|
||||
|
||||
if qrnn:
|
||||
print('Using QRNNs...')
|
||||
model_name = 'qrnn' if qrnn else 'lstm'
|
||||
lm_name = f'{model_name}_{pretrain_name}'
|
||||
pretrained_fname = (lm_name, f'itos_{pretrain_name}')
|
||||
class CLSHyperParams(LMHyperParams):
|
||||
# dir_path -> data/imdb/
|
||||
|
||||
ensure_paths_exists(data_dir,
|
||||
dataset_dir,
|
||||
model_dir,
|
||||
model_dir/f"{pretrained_fname[0]}.pth",
|
||||
model_dir/f"{pretrained_fname[1]}.pkl")
|
||||
def __post_init__(self, *args, **kwargs):
|
||||
super().__post_init__(*args, **kwargs)
|
||||
self.dataset_dir=self.dataset_path
|
||||
|
||||
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
|
||||
@property
|
||||
def need_fine_tune_lm(self): return not (self.model_dir/f"enc_best.pth").exists()
|
||||
|
||||
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)
|
||||
def train_cls(self, num_lm_epochs, unfreeze=True, bs=70):
|
||||
data_clas, data_lm = self.load_cls_data(bs)
|
||||
|
||||
if qrnn:
|
||||
emb_sz, nh, nl = 400, 1550, 3
|
||||
else:
|
||||
emb_sz, nh, nl = 400, 1150, 3
|
||||
if self.need_fine_tune_lm: self.train_lm(num_lm_epochs, data_lm=data_lm)
|
||||
learn = self.create_cls_learner(data_clas)
|
||||
|
||||
lm_enc_finetuned = f"{lm_name}_{dataset}_enc"
|
||||
if fine_tune and not (model_dir/f"{lm_enc_finetuned}.pth").exists():
|
||||
print('Fine-tuning the language model...', lm_enc_finetuned)
|
||||
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)
|
||||
if bidir:
|
||||
learn.metrics = [accuracy_fwd, accuracy_bwd]
|
||||
else:
|
||||
learn.metrics = [accuracy]
|
||||
try:
|
||||
learn.load('cls_last')
|
||||
print("Loading last classfier")
|
||||
except FileNotFoundError:
|
||||
learn.load_encoder("enc_best")
|
||||
|
||||
learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7))
|
||||
learn.unfreeze()
|
||||
if num_lm_epochs > 0: learn.fit_one_cycle(num_lm_epochs, 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:
|
||||
print(f"Loading classifier {model_name}_{name}")
|
||||
learn.load(f'{model_name}_{name}')
|
||||
|
||||
except FileNotFoundError:
|
||||
learn.load_encoder(lm_enc_finetuned)
|
||||
print("loading encoder")
|
||||
train = True
|
||||
|
||||
if train:
|
||||
learn.true_wd = False
|
||||
print("Starting classifier training")
|
||||
learn.fit_one_cycle(1, 5e-2, moms=(0.8, 0.7), wd=1e-7)
|
||||
if unfreeze:
|
||||
learn.freeze_to(-2)
|
||||
learn.fit_one_cycle(1, slice(5e-2 / (2.6 ** 4), 5e-2), moms=(0.8, 0.7), wd=1e-7)
|
||||
|
||||
learn.freeze_to(-2)
|
||||
learn.fit_one_cycle(1, slice(5e-2 / (2.6 ** 4), 5e-2), moms=(0.8, 0.7), wd=1e-7)
|
||||
learn.freeze_to(-3)
|
||||
learn.fit_one_cycle(1, slice(5e-4 / (2.6 ** 4), 5e-4), moms=(0.8, 0.7), wd=1e-7)
|
||||
|
||||
learn.freeze_to(-3)
|
||||
learn.fit_one_cycle(1, slice(5e-4 / (2.6 ** 4), 5e-4), moms=(0.8, 0.7), wd=1e-7)
|
||||
|
||||
learn.unfreeze()
|
||||
learn.fit_one_cycle(2, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
|
||||
learn.unfreeze()
|
||||
learn.fit_one_cycle(2, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
|
||||
|
||||
print(f"Saving models at {learn.path / learn.model_dir}")
|
||||
learn.save(f'{model_name}_{name}')
|
||||
learn.save('cls_last', with_opt=False)
|
||||
return learn
|
||||
|
||||
results['accuracy'] = learn.recorder.metrics[-1][0]
|
||||
return results
|
||||
def create_cls_learner(self, data_clas):
|
||||
fastai.text.learner.default_dropout['language'] = self.dps
|
||||
classifier_learner = bilm_text_classifier_learner if self.bidir else text_classifier_learner
|
||||
learn = classifier_learner(data_clas, bptt=self.bptt, pad_token=PAD_TOKEN_ID,
|
||||
path=self.model_dir.parent, model_dir=self.model_dir.name,
|
||||
qrnn=self.qrnn, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, drop_mult=self.drop_mult)
|
||||
|
||||
learn.metrics = [accuracy_fwd, accuracy_bwd] if self.bidir else [accuracy]
|
||||
return learn
|
||||
|
||||
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'
|
||||
if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists():
|
||||
print('Reading the data...')
|
||||
toks, lbls = read_clas_data(dataset_dir, dataset, lang)
|
||||
# create the vocabulary
|
||||
counter = Counter(word for example in toks[TRN]+toks[TST]+toks[VAL] for word in example)
|
||||
itos = [word for word, count in counter.most_common(n=max_vocab)]
|
||||
itos.insert(0, PAD)
|
||||
itos.insert(0, UNK)
|
||||
vocab = Vocab(itos)
|
||||
stoi = vocab.stoi
|
||||
with open(vocab_file, 'wb') as f:
|
||||
pickle.dump(vocab, f)
|
||||
|
||||
ids = {}
|
||||
def load_cls_data(self, bs):
|
||||
tmp_dir = self.cache_dir
|
||||
tmp_dir.mkdir(exist_ok=True)
|
||||
vocab_file = tmp_dir / f'vocab_{self.lang}.pkl'
|
||||
if not (tmp_dir / f'{TRN}_{self.lang}_ids.npy').exists():
|
||||
print('Reading the data...')
|
||||
toks, lbls = read_clas_data(self.dataset_dir, self.dataset_dir.name, self.lang)
|
||||
# create the vocabulary
|
||||
counter = Counter(word for example in toks[TRN] + toks[TST] + toks[VAL] for word in example)
|
||||
itos = [word for word, count in counter.most_common(n=self.max_vocab)]
|
||||
itos.insert(0, PAD)
|
||||
itos.insert(0, UNK)
|
||||
vocab = Vocab(itos)
|
||||
stoi = vocab.stoi
|
||||
with open(vocab_file, 'wb') as f:
|
||||
pickle.dump(vocab, f)
|
||||
ids = {}
|
||||
for split in [TRN, VAL, TST]:
|
||||
ids[split] = np.array([([stoi.get(w, stoi[UNK]) for w in s])
|
||||
for s in toks[split]])
|
||||
np.save(tmp_dir / f'{split}_{self.lang}_ids.npy', ids[split])
|
||||
np.save(tmp_dir / f'{split}_{self.lang}_lbl.npy', lbls[split])
|
||||
else:
|
||||
print('Loading the pickled data...')
|
||||
ids, lbls = {}, {}
|
||||
for split in [TRN, VAL, TST]:
|
||||
ids[split] = np.load(tmp_dir / f'{split}_{self.lang}_ids.npy')
|
||||
lbls[split] = np.load(tmp_dir / f'{split}_{self.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])}.')
|
||||
for split in [TRN, VAL, TST]:
|
||||
ids[split] = np.array([([stoi.get(w, stoi[UNK]) for w in s])
|
||||
for s in toks[split]])
|
||||
np.save(tmp_dir / f'{split}_{lang}_ids.npy', ids[split])
|
||||
np.save(tmp_dir / f'{split}_{lang}_lbl.npy', lbls[split])
|
||||
else:
|
||||
print('Loading the pickled data...')
|
||||
ids, lbls = {}, {}
|
||||
for split in [TRN, VAL, TST]:
|
||||
ids[split] = np.load(tmp_dir / f'{split}_{lang}_ids.npy')
|
||||
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] = np.array([np.array(e, dtype=np.int) for e in ids[split]])
|
||||
lbls[split] = np.array([np.array(e, dtype=np.int) for e in lbls[split]])
|
||||
data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=np.concatenate([ids[TRN],ids[TST]]),
|
||||
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],
|
||||
train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs, classes={l:l for l in lbls[TRN]})
|
||||
|
||||
print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
|
||||
return data_clas, data_lm
|
||||
ids[split] = np.array([np.array(e, dtype=np.int) for e in ids[split]])
|
||||
lbls[split] = np.array([np.array(e, dtype=np.int) for e in lbls[split]])
|
||||
data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=np.concatenate([ids[TRN], ids[TST]]),
|
||||
valid_ids=ids[VAL], bs=bs, bptt=self.bptt, lm_type=self.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],
|
||||
train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs, classes={l: l for l in lbls[TRN]})
|
||||
|
||||
print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
|
||||
return data_clas, data_lm
|
||||
|
||||
if __name__ == '__main__':
|
||||
fire.Fire(new_train_clas)
|
||||
fire.Fire(CLSHyperParams)
|
||||
|
||||
Reference in New Issue
Block a user