diff --git a/finetune.py b/finetune.py index 48906f9..f0e8dd5 100644 --- a/finetune.py +++ b/finetune.py @@ -704,7 +704,7 @@ def basemodel_or_finetunedmodel_choice(value): elif value == "Existing finetuned model": basemodel_or_finetunedmodel = False -def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=255995, progress=gr.Progress()): +def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=255995, progress=gr.Progress()): pfc_check_fail() if "No Models Available" in model_to_train: @@ -983,6 +983,8 @@ def train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, eval_samples=eval_samples, training_assets=training_assets, c_logger=c_logger, + warmup=warm_up, + ) if(disable_shared_memory): @@ -1789,6 +1791,11 @@ if __name__ == "__main__": label="Disable Shared Memory Use", scale=1, ) + warm_up = gr.Checkbox( + value=False, + label="Perform Small a Warmup Learning", + scale=1, + ) train_csv = gr.Textbox( label="Train CSV:", scale=2, @@ -1875,7 +1882,7 @@ if __name__ == "__main__": grad_acumm = gr.Slider( label="Grad accumulation steps:", minimum=1, - maximum=32, + maximum=128, step=1, value=args.grad_acumm, ) @@ -1906,7 +1913,7 @@ if __name__ == "__main__": demo.load(load_metrics, None, [model_data, train_time], every=1) demo.load(read_logs, None, logs_tts_train, every=1) - def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, progress=gr.Progress()): + def train_model(language, train_csv, eval_csv, learning_rates, model_to_train, num_epochs, batch_size, grad_acumm, max_audio_length, speaker_name_input_training, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, progress=gr.Progress()): clear_gpu_cache() global out_path if speaker_name_input_training and speaker_name_input_training != 'personsname': @@ -1914,8 +1921,8 @@ if __name__ == "__main__": if not train_csv or not eval_csv: if (out_path / "metadata_eval.csv").exists() and (out_path / "metadata_train.csv").exists(): - train_csv = out_path / "metadata_eval.csv" - eval_csv = out_path / "metadata_train.csv" + train_csv = out_path / "metadata_train.csv" + eval_csv = out_path / "metadata_eval.csv" print("Using existing metadata and training csv.") else: return "You need to run the data processing step or manually set `Train CSV` and `Eval CSV` fields !", "", "", "", "" @@ -1924,7 +1931,7 @@ if __name__ == "__main__": max_audio_length = int(max_audio_length * 22050) learning_rate = float(learning_rates) # Convert the learning rate value to a float progress(0, "Initializing training...") - config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, max_audio_length=max_audio_length, progress=gr.Progress()) + config_path, original_xtts_checkpoint, vocab_file, exp_path, speaker_wav = train_gpt(language, num_epochs, batch_size, grad_acumm, train_csv, eval_csv, learning_rate, model_to_train, continue_run, disable_shared_memory, learning_rate_scheduler, optimizer, warm_up, max_audio_length=max_audio_length, progress=gr.Progress()) # copy original files to avoid parameters changes issues shutil.copy(config_path, exp_path) @@ -2304,7 +2311,8 @@ if __name__ == "__main__": continue_run, disable_shared_memory, learning_rate_scheduler, - optimizer + optimizer, + warm_up ], outputs=[progress_train, xtts_config, xtts_vocab, xtts_checkpoint, speaker_reference_audio, speaker_name_input_testing], ) diff --git a/trainer_alltalk/WarmUpScheduler.py b/trainer_alltalk/WarmUpScheduler.py new file mode 100644 index 0000000..9406c82 --- /dev/null +++ b/trainer_alltalk/WarmUpScheduler.py @@ -0,0 +1,33 @@ +from torch.optim.lr_scheduler import LRScheduler + +class WarmUpScheduler(LRScheduler): + def __init__(self, optimizer, total_epochs, warmup_lr, target_lr, after_scheduler=None): + self.total_epochs = total_epochs + self.warmup_epochs = int(total_epochs * 0.1) + self.warmup_lr = warmup_lr + self.target_lr = target_lr + self.after_scheduler = after_scheduler + print(f"[WARMUP] {self.warmup_epochs} epochs") + super(WarmUpScheduler, self).__init__(optimizer) + + def get_lr(self): + if self.last_epoch < self.warmup_epochs: + progress = self.last_epoch / self.warmup_epochs + return [self.warmup_lr + progress * (self.target_lr - self.warmup_lr) for _ in self.base_lrs] + if self.after_scheduler: + return self.after_scheduler.get_last_lr() + return [base_lr for base_lr in self.base_lrs] + + def step(self, epoch=None): + if epoch is not None and epoch < self.warmup_epochs: + self.last_epoch = epoch + elif self.after_scheduler and epoch is not None: + if epoch < self.warmup_epochs: + self.last_epoch = epoch + else: + self.after_scheduler.step(epoch - self.warmup_epochs) + self.last_epoch = epoch + else: + super(WarmUpScheduler, self).step(epoch) + if self.after_scheduler and self.last_epoch >= self.warmup_epochs: + self.after_scheduler.step(epoch - self.warmup_epochs if epoch is not None else None) \ No newline at end of file diff --git a/trainer_alltalk/trainer.py b/trainer_alltalk/trainer.py index e786a2e..9b345ec 100644 --- a/trainer_alltalk/trainer.py +++ b/trainer_alltalk/trainer.py @@ -58,6 +58,8 @@ from trainer.utils.distributed import ( rank_zero_only, ) +from trainer_alltalk.WarmUpScheduler import WarmUpScheduler + logger = logging.getLogger("trainer") if is_apex_available(): @@ -309,6 +311,7 @@ class Trainer: parse_command_line_args: bool = True, callbacks: Dict[str, Callable] = {}, gpu: int = None, + warmup: bool = False, ) -> None: """Simple yet powerful πŸΈπŸ’¬ TTS trainer for PyTorch. It can train all the available `tts` and `vocoder` models or easily be customized. @@ -429,6 +432,7 @@ class Trainer: self.overfit_batch = args.overfit_batch self.skip_train_epoch = args.skip_train_epoch self.start_with_eval = args.start_with_eval + self.warmup = warmup assert self.grad_accum_steps > 0, " [!] grad_accum_steps must be greater than 0." @@ -569,6 +573,16 @@ class Trainer: self.scheduler, self.args, self.config, self.restore_epoch, self.restore_step ) + if self.warmup: + warmup_scheduler = WarmUpScheduler( + optimizer=self.optimizer, + total_epochs=self.config.epochs, + warmup_lr=1e-7, + target_lr=self.config.lr, + after_scheduler=self.scheduler + ) + self.scheduler = warmup_scheduler + # DISTRIBUTED if self.use_pt_ddp: self.model = DDP_th(self.model, device_ids=[args.rank], output_device=args.rank) @@ -1542,6 +1556,7 @@ class Trainer: self.model.train() torch.set_grad_enabled(True) + epoch_time = time.time() - epoch_start_time self.callbacks.on_train_epoch_end(self)