def main():
"""Run training process."""
parser = argparse.ArgumentParser(description="Train Tacotron2")
parser.add_argument("--outdir", type=str, required=True, help="directory to save checkpoints.")
parser.add_argument("--rootdir", type=str, required=True, help="dataset directory root")
parser.add_argument("--resume",default="",type=str,nargs="?",help='checkpoint file path to resume training. (default="")')
parser.add_argument("--verbose",type=int,default=1,help="logging level. higher is more logging. (default=1)")
parser.add_argument("--batch-size", default=12, type=int, help="batch size.")
parser.add_argument("--mixed_precision",default=0,type=int,help="using mixed precision for generator or not.")
args = parser.parse_args()
if args.resume is not None and os.path.isdir(args.resume):
args.resume = tf.train.latest_checkpoint(args.resume)
# set mixed precision config
if args.mixed_precision == 1:
tf.config.optimizer.set_experimental_options({"auto_mixed_precision": True})
args.mixed_precision = bool(args.mixed_precision)
# set logger
log_format = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
if args.verbose > 1:
logging.basicConfig(level=logging.DEBUG,stream=sys.stdout,format=log_format)
elif args.verbose > 0:
logging.basicConfig(level=logging.INFO,stream=sys.stdout,format=log_format)
else:
logging.basicConfig(level=logging.WARN,stream=sys.stdout,format=log_format)
logging.warning("Skip DEBUG/INFO messages")
# check directory existence(checkpoint)
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
# select processor
Processor = JSpeechProcessor # for test
processor = Processor(rootdir=args.rootdir)
config = Config(args.outdir, args.batch_size, processor.vocab_size())
max_mel_length = processor.max_feat_length() // config.n_mels
max_seq_length = processor.max_seq_length()
# split train and test
train_split, valid_split = train_test_split(processor.items, test_size=config.test_size,random_state=42,shuffle=True)
train_dataset = generate_datasets(train_split, config, max_mel_length, max_seq_length)
valid_dataset = generate_datasets(valid_split, config, max_mel_length, max_seq_length)
# define trainer
trainer = Tacotron2Trainer(
config=config,
strategy=STRATEGY,
steps=0,
epochs=0,
is_mixed_precision=args.mixed_precision
)
with STRATEGY.scope():
# define model.
tacotron2 = TFTacotron2(config=config, training=True, name="tacotron2")
#build
input_ids = np.array([[1, 2, 3, 4, 5, 6, 7, 8, 9]])
input_lengths = np.array([9])
speaker_ids = np.array([0])
mel_outputs = np.random.normal(size=(1, 50, config.n_mels)).astype(np.float32)
mel_lengths = np.array([50])
tacotron2(input_ids,input_lengths,speaker_ids,mel_outputs,mel_lengths,10,training=True)
tacotron2.summary()
# AdamW for tacotron2
learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
initial_learning_rate=config.initial_learning_rate,
decay_steps=config.decay_steps,
end_learning_rate=config.end_learning_rate,
)
learning_rate_fn = WarmUp(
initial_learning_rate=config.initial_learning_rate,
decay_schedule_fn=learning_rate_fn,
warmup_steps=int(config.train_max_steps* config.warmup_proportion),
)
optimizer = AdamWeightDecay(
learning_rate=learning_rate_fn,
weight_decay_rate=config.weight_decay,
beta_1=0.9,
beta_2=0.98,
epsilon=1e-6,
exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
)
_ = optimizer.iterations
# compile trainer
trainer.compile(model=tacotron2, optimizer=optimizer)
# start training
try:
trainer.fit(train_dataset,valid_dataset,saved_path=os.path.join(args.outdir, "checkpoints/"),resume=args.resume)
except KeyboardInterrupt:
trainer.save_checkpoint()
logging.info(f"Successfully saved checkpoint @ {trainer.steps}steps.")
# initialize melgan model
with open(config_lp.multiband_melgan_baker) as f:
melgan_config = yaml.load(f, Loader=yaml.Loader)
melgan_config = MelGANGeneratorConfig(
**melgan_config["multiband_melgan_generator_params"])
melgan = TFMelGANGenerator(config=melgan_config, name='mb_melgan')
melgan._build()
melgan.load_weights(config_lp.multiband_melgan_pretrained_path)
# initialize Tacotron2 model.
with open(config_lp.tacotron2_baker) as f:
config = yaml.load(f, Loader=yaml.Loader)
config = Tacotron2Config(**config["tacotron2_params"])
tacotron2 = TFTacotron2(config=config,
training=False,
name="tacotron2v2",
enable_tflite_convertible=True)
# Newly added :
tacotron2.setup_window(win_front=6, win_back=6)
tacotron2.setup_maximum_iterations(3000)
tacotron2._build()
tacotron2.load_weights(config_lp.tacotron2_pretrained_path)
tacotron2.summary()
# Concrete Function
tacotron2_concrete_function = tacotron2.inference_tflite.get_concrete_function(
)
converter = tf.lite.TFLiteConverter.from_concrete_functions(
def main():
"""Running extract tacotron-2 durations."""
parser = argparse.ArgumentParser(
description="Extract durations from charactor with trained Tacotron-2 "
"(See detail in tensorflow_tts/example/tacotron-2/extract_duration.py).")
parser.add_argument("--rootdir",default=None,type=str, required=True,help="directory including ids/durations files.",)
parser.add_argument("--outdir", type=str, required=True, help="directory to save generated speech.")
parser.add_argument("--checkpoint", type=str, required=True, help="checkpoint file to be loaded." )
parser.add_argument("--use-norm", default=1, type=int, help="usr norm-mels for train or raw.")
parser.add_argument("--batch-size", default=8, type=int, help="batch size.")
parser.add_argument("--win-front", default=3, type=int, help="win-front.")
parser.add_argument("--win-back", default=3, type=int, help="win-front.")
parser.add_argument("--use-window-mask", default=1, type=int, help="toggle window masking." )
parser.add_argument("--save-alignment", default=1, type=int, help="save-alignment.")
parser.add_argument("--dataset_mapping", default="dump/baker_mapper.json", type=str, )
parser.add_argument("--config", default=None, type=str, required=True, help="yaml format configuration file. if not explicitly provided, it will be searched in the checkpoint directory. (default=None)", )
parser.add_argument( "--verbose", type=int,default=1, help="logging level. higher is more logging. (default=1)", )
args = parser.parse_args()
print(args)
# set logger
if args.verbose > 1:
logging.basicConfig(level=logging.DEBUG,format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",)
elif args.verbose > 0:
logging.basicConfig(level=logging.INFO,format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",)
else:
logging.basicConfig(level=logging.WARN,format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",)
logging.warning("Skip DEBUG/INFO messages")
# check directory existence
if not os.path.exists(args.outdir):
os.makedirs(args.outdir)
if not os.path.exists(args.outdir+'_align_fig'):
os.makedirs(args.outdir+'_align_fig')
# load config
with open(args.config) as f:
config = yaml.load(f, Loader=yaml.Loader)
config.update(vars(args))
if config["format"] == "npy":
char_query = "*-ids.npy"
mel_query = "*-raw-feats.npy" if args.use_norm is False else "*-norm-feats.npy"
char_load_fn = np.load
mel_load_fn = np.load
else:
raise ValueError("Only npy is supported.")
with open(args.dataset_mapping) as f:
dataset_mapping = json.load(f)
speakers_map = dataset_mapping["speakers_map"]
# Check n_speakers matches number of speakers in speakers_map
n_speakers = config["tacotron2_params"]["n_speakers"]
# define data-loader
dataset = CharactorMelDataset(
dataset=config["tacotron2_params"]["dataset"],
root_dir=args.rootdir,
charactor_query=char_query,
mel_query=mel_query,
charactor_load_fn=char_load_fn,
mel_load_fn=mel_load_fn,
reduction_factor=config["tacotron2_params"]["reduction_factor"],
use_fixed_shapes=True,
speakers_map=speakers_map,
)
dataset = dataset.create(allow_cache=True, batch_size=args.batch_size, drop_remainder=False)
# define model and load checkpoint
tacotron2 = TFTacotron2(config=Tacotron2Config(**config["tacotron2_params"]), name="tacotron2",)
tacotron2._build() # build model to be able load_weights.
tacotron2.load_weights(args.checkpoint)
# apply tf.function for tacotron2.
tacotron2 = tf.function(tacotron2, experimental_relax_shapes=True)
for data in tqdm(dataset, desc="[Extract Duration]"):
utt_ids = data["utt_ids"]
input_lengths = data["input_lengths"]
mel_lengths = data["mel_lengths"]
utt_ids = utt_ids.numpy()
real_mel_lengths = data["real_mel_lengths"]
del data["real_mel_lengths"]
# tacotron2 inference.
mel_outputs, post_mel_outputs, stop_outputs, alignment_historys = tacotron2(**data,use_window_mask=args.use_window_mask,win_front=args.win_front,win_back=args.win_back,training=True,)
# convert to numpy
alignment_historys = alignment_historys.numpy()
for i, alignment in enumerate(alignment_historys):
real_char_length = input_lengths[i].numpy()
real_mel_length = real_mel_lengths[i].numpy()
alignment_mel_length = int(np.ceil( real_mel_length / config["tacotron2_params"]["reduction_factor"]) )
alignment = alignment[:real_char_length, :alignment_mel_length]
d = get_duration_from_alignment(alignment) # [max_char_len]
d = d * config["tacotron2_params"]["reduction_factor"]
assert ( np.sum(d) >= real_mel_length ), f"{d}, {np.sum(d)}, {alignment_mel_length}, {real_mel_length}"
if np.sum(d) > real_mel_length:
rest = np.sum(d) - real_mel_length
# print(d, np.sum(d), real_mel_length)
if d[-1] > rest:
d[-1] -= rest
elif d[0] > rest:
d[0] -= rest
else:
d[-1] -= rest // 2
d[0] -= rest - rest // 2
assert d[-1] >= 0 and d[0] >= 0, f"{d}, {np.sum(d)}, {real_mel_length}"
saved_name = utt_ids[i].decode("utf-8")
# check a length compatible
assert ( len(d) == real_char_length ), f"different between len_char and len_durations, {len(d)} and {real_char_length}"
assert ( np.sum(d) == real_mel_length ), f"different between sum_durations and len_mel, {np.sum(d)} and {real_mel_length}"
# save D to folder.
np.save( os.path.join(args.outdir, f"{saved_name}-durations.npy"), d.astype(np.int32), allow_pickle=False,)
# save alignment to debug.
if args.save_alignment == 1:
figname = os.path.join(args.outdir+'_align_fig/', f"{saved_name}_alignment.png")
#print(figname)
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111)
ax.set_title(f"Alignment of {saved_name}")
im = ax.imshow( alignment, aspect="auto", origin="lower", interpolation="none" )
fig.colorbar(im, ax=ax)
xlabel = "Decoder timestep"
plt.xlabel(xlabel)
plt.ylabel("Encoder timestep")
plt.tight_layout()
plt.savefig(figname)
plt.close()