add warmup

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)