-
Notifications
You must be signed in to change notification settings - Fork 132
/
tacotron_train.py
65 lines (53 loc) · 2.75 KB
/
tacotron_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
61
62
63
64
65
import argparse
import os
#os.environ["CUDA_VISIBLE_DEVICES"] = ""
from time import sleep
from tacotron.utils import infolog
import tensorflow as tf
from tacotron_hparams import hparams
from tacotron_hparams import hparams as hp
from tacotron.utils.infolog import log
from tacotron.synthesize import tacotron_synthesize
from tacotron.train import tacotron_train
import sys
log = infolog.log
def prepare_run(args):
modified_hp = hparams.parse(args.hparams)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(args.tf_log_level)
run_name = args.name or args.model
log_dir = os.path.join(args.base_dir, 'logs-{}'.format(run_name))
os.makedirs(log_dir, exist_ok=True)
infolog.init(os.path.join(log_dir, 'Terminal_train_log'), run_name)
return log_dir, modified_hp
def train(args, log_dir, hparams):
log('\n#############################################################\n')
log('Tacotron Train\n')
log('###########################################################\n')
tacotron_train(args, log_dir, hparams)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--data_dir', default=hp.dataset)
parser.add_argument('--base_dir', default=os.getcwd())
parser.add_argument('--hparams', default='', help='Hyperparameter overrides as a comma-separated list of name=value pairs')
parser.add_argument('--tacotron_input', default=hp.tacotron_input)
parser.add_argument('--name', help='Name of logging directory.')
parser.add_argument('--model', default='Tacotron-2')
parser.add_argument('--input_dir', default=hp.feat_out_dir, help='folder to contain inputs sentences/targets')
parser.add_argument('--output_dir', default='output', help='folder to contain synthesized mel spectrograms')
parser.add_argument('--restore', type=bool, default=True, help='Set this to False to do a fresh training')
parser.add_argument('--summary_interval', type=int, default=1000,
help='Steps between running summary ops')
parser.add_argument('--embedding_interval', type=int, default=10000,
help='Steps between updating embeddings projection visualization')
parser.add_argument('--checkpoint_interval', type=int, default=500,
help='Steps between writing checkpoints')
parser.add_argument('--eval_interval', type=int, default=5000,
help='Steps between eval on test data')
parser.add_argument('--tacotron_train_steps', type=int, default=300000, help='total number of tacotron training steps')
parser.add_argument('--tf_log_level', type=int, default=1, help='Tensorflow C++ log level.')
args = parser.parse_args()
assert args.model == 'Tacotron-2', 'args.model != Tacotron-2'
log_dir, hparams = prepare_run(args)
train(args, log_dir, hparams)
if __name__ == '__main__':
main()