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

@@ -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):

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())

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())