def load_model(self, checkpoint_path: str = None, verbose=True):
model = self.get_model()
self.compile_model(model)
ckpt = tf.train.Checkpoint(net=model)
manager = tf.train.CheckpointManager(ckpt,
self.weights_dir,
max_to_keep=None)
if checkpoint_path:
ckpt.restore(checkpoint_path)
if verbose:
print(
f'restored weights from {checkpoint_path} at step {model.step}'
)
else:
if manager.latest_checkpoint is None:
print(
f'WARNING: could not find weights file. Trying to load from \n {self.weights_dir}.'
)
print(
'Edit data_config.yaml to point at the right log directory.'
)
ckpt.restore(manager.latest_checkpoint)
if verbose:
print(
f'restored weights from {manager.latest_checkpoint} at step {model.step}'
)
decoder_prenet_dropout = piecewise_linear_schedule(
model.step, self.config['decoder_prenet_dropout_schedule'])
reduction_factor = None
if self.model_kind == 'autoregressive':
reduction_factor = reduction_schedule(
model.step, self.config['reduction_factor_schedule'])
model.set_constants(reduction_factor=reduction_factor,
decoder_prenet_dropout=decoder_prenet_dropout)
return model
if config['debug'] is True:
print('\nWARNING: DEBUG is set to True. Training in eager mode.')
# main event
print('\nTRAINING')
losses = []
_ = train_dataset.next_batch()
t = trange(model.step, config['max_steps'], leave=True)
for _ in t:
t.set_description(f'step {model.step}')
mel, phonemes, stop = train_dataset.next_batch()
decoder_prenet_dropout = piecewise_linear_schedule(
model.step, config['decoder_prenet_dropout_schedule'])
learning_rate = piecewise_linear_schedule(model.step,
config['learning_rate_schedule'])
reduction_factor = reduction_schedule(model.step,
config['reduction_factor_schedule'])
drop_n_heads = tf.cast(
reduction_schedule(model.step, config['head_drop_schedule']), tf.int32)
t.display(f'reduction factor {reduction_factor}', pos=10)
model.set_constants(decoder_prenet_dropout=decoder_prenet_dropout,
learning_rate=learning_rate,
reduction_factor=reduction_factor,
drop_n_heads=drop_n_heads)
output = model.train_step(inp=phonemes, tar=mel, stop_prob=stop)
losses.append(float(output['loss']))
t.display(f'step loss: {losses[-1]}', pos=1)
for pos, n_steps in enumerate(config['n_steps_avg_losses']):
if len(losses) > n_steps:
t.display(
f'{n_steps}-steps average loss: {sum(losses[-n_steps:]) / n_steps}',
else:
print(f'\nstarting training from scratch')
if config['debug'] is True:
print('\nWARNING: DEBUG is set to True. Training in eager mode.')
# main event
print('\nTRAINING')
losses = []
_ = train_dataset.next_batch()
t = trange(model.step, config['max_steps'], leave=True)
for _ in t:
t.set_description(f'step {model.step}')
mel, phonemes, stop = train_dataset.next_batch()
decoder_prenet_dropout = piecewise_linear_schedule(model.step, config['decoder_prenet_dropout_schedule'])
learning_rate = piecewise_linear_schedule(model.step, config['learning_rate_schedule'])
reduction_factor = reduction_schedule(model.step, config['reduction_factor_schedule'])
drop_n_heads = tf.cast(reduction_schedule(model.step, config['head_drop_schedule']), tf.int32)
t.display(f'reduction factor {reduction_factor}', pos=10)
model.set_constants(decoder_prenet_dropout=decoder_prenet_dropout,
learning_rate=learning_rate,
reduction_factor=reduction_factor,
drop_n_heads=drop_n_heads)
output = model.train_step(inp=phonemes,
tar=mel,
stop_prob=stop)
losses.append(float(output['loss']))
t.display(f'step loss: {losses[-1]}', pos=1)
for pos, n_steps in enumerate(config['n_steps_avg_losses']):
if len(losses) > n_steps:
t.display(f'{n_steps}-steps average loss: {sum(losses[-n_steps:]) / n_steps}', pos=pos + 2)
else:
print(f'\nstarting training from scratch')
# main event
print('\nTRAINING')
losses = []
test_batch = val_dataset.next_batch()
t = trange(model.step, config['max_steps'], leave=True)
for _ in t:
t.set_description(f'step {model.step}')
mel, phonemes, durations = train_dataset.next_batch()
learning_rate = piecewise_linear_schedule(model.step,
config['learning_rate_schedule'])
decoder_prenet_dropout = piecewise_linear_schedule(
model.step, config['decoder_prenet_dropout_schedule'])
drop_n_heads = tf.cast(
reduction_schedule(model.step, config['head_drop_schedule']), tf.int32)
model.set_constants(decoder_prenet_dropout=decoder_prenet_dropout,
learning_rate=learning_rate,
drop_n_heads=drop_n_heads)
output = model.train_step(input_sequence=phonemes,
target_sequence=mel,
target_durations=durations)
losses.append(float(output['loss']))
t.display(f'step loss: {losses[-1]}', pos=1)
for pos, n_steps in enumerate(config['n_steps_avg_losses']):
if len(losses) > n_steps:
t.display(
f'{n_steps}-steps average loss: {sum(losses[-n_steps:]) / n_steps}',
pos=pos + 2)
if config_dict['debug'] is True:
print('\nWARNING: DEBUG is set to True. Training in eager mode.')
# main event
print('\nTRAINING')
losses = []
test_mel, test_phonemes, test_durs, test_fname = valid_dataset.next_batch()
t = trange(model.step, config_dict['max_steps'], leave=True)
for _ in t:
t.set_description(f'step {model.step}')
mel, phonemes, durations, fname = train_dataset.next_batch()
learning_rate = piecewise_linear_schedule(
model.step, config_dict['learning_rate_schedule'])
decoder_prenet_dropout = piecewise_linear_schedule(
model.step, config_dict['decoder_prenet_dropout_schedule'])
drop_n_heads = tf.cast(
reduction_schedule(model.step, config_dict['head_drop_schedule']),
tf.int32)
model.set_constants(decoder_prenet_dropout=decoder_prenet_dropout,
learning_rate=learning_rate,
drop_n_heads=drop_n_heads)
output = model.train_step(input_sequence=phonemes,
target_sequence=mel,
target_durations=durations)
losses.append(float(output['loss']))
t.display(f'step loss: {losses[-1]}', pos=1)
for pos, n_steps in enumerate(config_dict['n_steps_avg_losses']):
if len(losses) > n_steps:
t.display(
f'{n_steps}-steps average loss: {sum(losses[-n_steps:]) / n_steps}',
pos=pos + 2)
