Merge pull request #26 from manmay-nakhashi/main

added aligner loss

Author: lucidrains

naturalspeech2_pytorch/naturalspeech2_pytorch.py

@@ -28,7 +28,7 @@
 from beartype.door import is_bearable
 
 from naturalspeech2_pytorch.attend import Attend
-from naturalspeech2_pytorch.aligner import Aligner
+from naturalspeech2_pytorch.aligner import Aligner, ForwardSumLoss
 from naturalspeech2_pytorch.utils.tokenizer import Tokenizer, ESpeak
 from naturalspeech2_pytorch.utils.utils import average_over_durations, create_mask
 from naturalspeech2_pytorch.version import __version__
@@ -1118,7 +1118,8 @@ def __init__(
         audio_to_mel_kwargs: dict = dict(),
         scale = 1., # this will be set to < 1. for better convergence when training on higher resolution images
         duration_loss_weight = 1.,
-        pitch_loss_weight = 1.
+        pitch_loss_weight = 1.,
+        aligner_loss_weight = 1.
     ):
         super().__init__()
 
@@ -1154,6 +1155,7 @@ def __init__(
             self.duration_pitch = DurationPitchPredictor(dim=duration_pitch_dim)
             self.aligner = Aligner(dim_in=aligner_dim_in, dim_hidden=aligner_dim_hidden, attn_channels=aligner_attn_channels)
             self.pitch_emb = nn.Embedding(pitch_emb_dim, pitch_emb_pp_hidden_dim)
+            self.aligner_loss = ForwardSumLoss()
 
         # rest of ddpm
 
@@ -1207,6 +1209,7 @@ def __init__(
 
         self.duration_loss_weight = duration_loss_weight
         self.pitch_loss_weight = pitch_loss_weight
+        self.aligner_loss_weight = aligner_loss_weight
 
     @property
     def device(self):
@@ -1488,10 +1491,12 @@ def forward(
 
             pitch = rearrange(pitch, 'b 1 d -> b d')
             pitch_loss = F.l1_loss(pitch, pitch_pred)
-
+            align_loss = self.aligner_loss(aln_log , text_lens, mel_lens)
             # weigh the losses
 
-            aux_loss = duration_loss * self.duration_loss_weight + pitch_loss + self.pitch_loss_weight
+            aux_loss = (duration_loss * self.duration_loss_weight) \
+                    + (pitch_loss * self.pitch_loss_weight) \
+                    + (align_loss * self.aligner_loss_weight)
 
         # automatically encode raw audio to residual vq with codec