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)