-
Notifications
You must be signed in to change notification settings - Fork 12
/
t2t-trainer
36 lines (29 loc) · 846 Bytes
/
t2t-trainer
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
#!/usr/bin/env python
"""Trainer for Tensor2Tensor.
This script is used to train your models in Tensor2Tensor.
For example, to train a shake-shake model on MNIST run this:
t2t-trainer \
--generate_data \
--problem=image_mnist \
--data_dir=~/t2t_data \
--tmp_dir=~/t2t_data/tmp
--model=shake_shake \
--hparams_set=shake_shake_quick \
--output_dir=~/t2t_train/mnist1 \
--train_steps=1000 \
--eval_steps=100
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from tensor2tensor.bin import t2t_trainer
import tensorflow as tf
# import warnings
# warnings.filterwarnings("ignore")
# import os
# os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
def main(argv):
t2t_trainer.main(argv)
if __name__ == "__main__":
tf.logging.set_verbosity(tf.logging.ERROR)
tf.app.run()