-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
45 lines (36 loc) · 1.32 KB
/
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
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
# Imports
import tensorflow as tf
import lib
from tfmodel import model_input_fn
def main(args):
for _ in range(lib._EPOCHS // lib._EPOCHS_PER_EVAL):
tensors_to_log = {
'cross_entropy': 'cross_entropy'
}
logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=1)
train_hooks = [logging_hook]
# cfg = tf.estimator.RunConfig().replace(save_checkpoints_secs=1e9)
run_config = tf.estimator.RunConfig(
model_dir=lib.model_dir
)
# Create the Estimator
model = tf.estimator.Estimator(
model_fn=model_input_fn,
model_dir=lib.model_dir,
config=run_config,
params={
'n_classes': lib._NUM_CLASSES,
'batch_size': lib._BATCH_SIZE
}
)
tf.logging.info("Start Evaluation")
model.train(
input_fn=lambda: lib.input_fn_from_tfrecords(True, data_dir=lib.records_dir, batch_size=lib._BATCH_SIZE, epochs=lib._EPOCHS),
hooks=train_hooks
)
if __name__ == '__main__':
tf.logging.set_verbosity(tf.logging.INFO)
tf.app.run(main=main)