wassname/multifit
1
0
Fork 0
mirror of https://github.com/wassname/multifit.git synced 2026-09-09 11:27:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add support for QRNN training to train_clas

Browse Source
This commit is contained in:
Piotr Czapla
2018-11-16 17:04:52 +01:00
parent d4698e095c
commit f0e538f621
1 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ulmfit/train_clas.py
+7 -4
View File
@@ -52,7 +52,13 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
f'Error: Name of data directory should be data, not {data_dir.name}.'
dataset_dir = data_dir / dataset
model_dir = Path(model_dir)
pretrained_fname = (f'lstm_{pretrain_name}', f'itos_{pretrain_name}')
if qrnn:
print('Using QRNNs...')
model_name = 'qrnn' if qrnn else 'lstm'
pretrained_fname = (f'{model_name}_{pretrain_name}', f'itos_{pretrain_name}')
ensure_paths_exists(data_dir,
dataset_dir,
@@ -60,9 +66,6 @@ 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")
if qrnn:
print('Using QRNNs...')
model_name = 'qrnn' if qrnn else 'lstm'
tmp_dir = dataset_dir / 'tmp'
tmp_dir.mkdir(exist_ok=True)