From 34b80e0280ae3396714f245bf4ca33a8b5d67baf Mon Sep 17 00:00:00 2001
From: a-froghyar <adamfroghyar@gmail.com>
Date: Tue, 12 Jul 2022 15:00:53 +0200
Subject: [PATCH] feat: updated recipes and lr fix (#1718)

- updated the recipes activating more losses for more stable training
- re-enabling guided attention loss
- fixed a bug about not the correct lr fetched for logging
---
 TTS/utils/capacitron_optimizer.py                    |  2 ++
 .../tacotron1-Capacitron/train_capacitron_t1.py      | 12 +++++-------
 .../tacotron2-Capacitron/train_capacitron_t2.py      | 12 +++++-------
 3 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/TTS/utils/capacitron_optimizer.py b/TTS/utils/capacitron_optimizer.py
index c9f075af..fac7d8a0 100644
--- a/TTS/utils/capacitron_optimizer.py
+++ b/TTS/utils/capacitron_optimizer.py
@@ -34,6 +34,8 @@ class CapacitronOptimizer:
         self.primary_optimizer.zero_grad()
 
     def step(self):
+        # Update param groups to display the correct learning rate
+        self.param_groups = self.primary_optimizer.param_groups
         self.primary_optimizer.step()
 
     def zero_grad(self):
diff --git a/recipes/blizzard2013/tacotron1-Capacitron/train_capacitron_t1.py b/recipes/blizzard2013/tacotron1-Capacitron/train_capacitron_t1.py
index 52c6098f..060cb9d4 100644
--- a/recipes/blizzard2013/tacotron1-Capacitron/train_capacitron_t1.py
+++ b/recipes/blizzard2013/tacotron1-Capacitron/train_capacitron_t1.py
@@ -48,7 +48,6 @@ config = TacotronConfig(
     precompute_num_workers=24,
     run_eval=True,
     test_delay_epochs=5,
-    ga_alpha=0.0,
     r=2,
     optimizer="CapacitronOptimizer",
     optimizer_params={"RAdam": {"betas": [0.9, 0.998], "weight_decay": 1e-6}, "SGD": {"lr": 1e-5, "momentum": 0.9}},
@@ -68,16 +67,15 @@ config = TacotronConfig(
     datasets=[dataset_config],
     lr=1e-3,
     lr_scheduler="StepwiseGradualLR",
-    lr_scheduler_params={"gradual_learning_rates": [[0, 1e-3], [2e4, 5e-4], [4e5, 3e-4], [6e4, 1e-4], [8e4, 5e-5]]},
+    lr_scheduler_params={"gradual_learning_rates": [[0, 1e-3], [2e4, 5e-4], [4e4, 3e-4], [6e4, 1e-4], [8e4, 5e-5]]},
     scheduler_after_epoch=False,  # scheduler doesn't work without this flag
-    # Need to experiment with these below for capacitron
     loss_masking=False,
     decoder_loss_alpha=1.0,
     postnet_loss_alpha=1.0,
-    postnet_diff_spec_alpha=0.0,
-    decoder_diff_spec_alpha=0.0,
-    decoder_ssim_alpha=0.0,
-    postnet_ssim_alpha=0.0,
+    postnet_diff_spec_alpha=1.0,
+    decoder_diff_spec_alpha=1.0,
+    decoder_ssim_alpha=1.0,
+    postnet_ssim_alpha=1.0,
 )
 
 ap = AudioProcessor(**config.audio.to_dict())
diff --git a/recipes/blizzard2013/tacotron2-Capacitron/train_capacitron_t2.py b/recipes/blizzard2013/tacotron2-Capacitron/train_capacitron_t2.py
index 9d588038..1bd2a036 100644
--- a/recipes/blizzard2013/tacotron2-Capacitron/train_capacitron_t2.py
+++ b/recipes/blizzard2013/tacotron2-Capacitron/train_capacitron_t2.py
@@ -52,7 +52,6 @@ config = Tacotron2Config(
     precompute_num_workers=24,
     run_eval=True,
     test_delay_epochs=5,
-    ga_alpha=0.0,
     r=2,
     optimizer="CapacitronOptimizer",
     optimizer_params={"RAdam": {"betas": [0.9, 0.998], "weight_decay": 1e-6}, "SGD": {"lr": 1e-5, "momentum": 0.9}},
@@ -77,21 +76,20 @@ config = Tacotron2Config(
         "gradual_learning_rates": [
             [0, 1e-3],
             [2e4, 5e-4],
-            [4e5, 3e-4],
+            [4e4, 3e-4],
             [6e4, 1e-4],
             [8e4, 5e-5],
         ]
     },
     scheduler_after_epoch=False,  # scheduler doesn't work without this flag
     seq_len_norm=True,
-    # Need to experiment with these below for capacitron
     loss_masking=False,
     decoder_loss_alpha=1.0,
     postnet_loss_alpha=1.0,
-    postnet_diff_spec_alpha=0.0,
-    decoder_diff_spec_alpha=0.0,
-    decoder_ssim_alpha=0.0,
-    postnet_ssim_alpha=0.0,
+    postnet_diff_spec_alpha=1.0,
+    decoder_diff_spec_alpha=1.0,
+    decoder_ssim_alpha=1.0,
+    postnet_ssim_alpha=1.0,
 )
 
 ap = AudioProcessor(**config.audio.to_dict())