forked from okuchaiev/f-lm
-
Notifications
You must be signed in to change notification settings - Fork 2
/
single_lm_train.py
62 lines (49 loc) · 2.24 KB
/
single_lm_train.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# -*- coding: utf-8 -*-
"""
Entry point for training and eval
"""
import os
import tensorflow as tf
import predict
import prediction
from data_utils import Vocabulary, Dataset
from language_model import LM
from prediction import sentence_ppl
from run_utils import run_train, run_eval
tf.flags.DEFINE_string("logdir", "lm1b", "Logging directory.")
tf.flags.DEFINE_string("datadir", None, "Logging directory.")
tf.flags.DEFINE_string("mode", "train", "Whether to run 'train' or 'eval' model.")
tf.flags.DEFINE_string("hpconfig", "", "Overrides default hyper-parameters.")
tf.flags.DEFINE_integer("num_gpus", 8, "Number of GPUs used.")
tf.flags.DEFINE_integer("eval_steps", 70, "Number of eval steps.")
FLAGS = tf.flags.FLAGS
def main(_):
"""
Start either train or eval. Note hardcoded parts of path for training and eval data
"""
hps = LM.get_default_hparams().parse(FLAGS.hpconfig)
hps._set("num_gpus", FLAGS.num_gpus)
print '*****HYPER PARAMETERS*****'
print hps
print '**************************'
vocab = Vocabulary.from_file(os.path.join(FLAGS.datadir, "vocabulary.txt"))
if FLAGS.mode == "train":
#hps.batch_size = 256
dataset = Dataset(vocab, os.path.join(FLAGS.datadir, "train/train.txt"))
run_train(dataset, hps, os.path.join(FLAGS.logdir, "train"), ps_device="/gpu:0")
elif FLAGS.mode.startswith("eval_"):
print 'start eval'
data_dir = os.path.join(FLAGS.datadir, "eval/valid.txt")
#predict_model = prediction.Model("/Users/ruiyangwang/Desktop/f-lm/logs/test/train/model.ckpt-0", os.path.join(FLAGS.datadir, "vocabulary.txt"), hps)
dataset = Dataset(vocab, data_dir, deterministic=True)
#prefix_words = 'i like music very much'
#predict_model.predict_top(prefix_words)
#prediction.topkwords(prefix_words, dataset, hps, FLAGS.logdir, FLAGS.mode)
# sentence_ppl(prefix_words,dataset, hps, FLAGS.logdir, FLAGS.mode)
#print vocab
run_eval(dataset, hps, FLAGS.logdir, FLAGS.mode, FLAGS.eval_steps)
#model = predict.Model(hps, FLAGS.logdir, FLAGS.datadir)
#print model.predictnextkwords(prefix_words.split(), 5)
#print model.predictnextkwords('know'.split(), 5)
if __name__ == "__main__":
tf.app.run()