add warmup

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"
+                            train_csv = out_path / "metadata_train.csv"
-                            eval_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],
             )

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)

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)