Add masked MSE loss

loss_function.py

@@ -1,19 +1,27 @@
 from torch import nn
+import torch.nn.functional as F
 
 
 class Tacotron2Loss(nn.Module):
     def __init__(self):
         super(Tacotron2Loss, self).__init__()
 
-    def forward(self, model_output, targets):
+    @staticmethod
+    def masked_l2_loss(out, target, lengths):
+        num_not_padded = lengths.sum() * out.size(1)
+        loss = F.mse_loss(out, target, reduction="sum")
+        loss = loss / num_not_padded
+        return loss
+
+    def forward(self, model_output, targets, output_lengths):
         mel_target, gate_target = targets[0], targets[1]
         mel_target.requires_grad = False
         gate_target.requires_grad = False
         gate_target = gate_target.view(-1, 1)
 
         mel_out, mel_out_postnet, gate_out, _ = model_output
         gate_out = gate_out.view(-1, 1)
-        mel_loss = nn.MSELoss()(mel_out, mel_target) + \
-            nn.MSELoss()(mel_out_postnet, mel_target)
+        mel_loss = self.masked_l2_loss(mel_out, mel_target, output_lengths) + \
+            self.masked_l2_loss(mel_out_postnet, mel_target, output_lengths)
         gate_loss = nn.BCEWithLogitsLoss()(gate_out, gate_target)
         return mel_loss + gate_loss

train.py

@@ -132,7 +132,7 @@ def validate(model, criterion, valset, iteration, batch_size, n_gpus,
         for i, batch in enumerate(val_loader):
             x, y = model.parse_batch(batch)
             y_pred = model(x)
-            loss = criterion(y_pred, y)
+            loss = criterion(y_pred, y, x[-1])
             if distributed_run:
                 reduced_val_loss = reduce_tensor(loss.data, n_gpus).item()
             else:
@@ -214,7 +214,7 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
             x, y = model.parse_batch(batch)
             y_pred = model(x)
 
-            loss = criterion(y_pred, y)
+            loss = criterion(y_pred, y, x[-1])
             if hparams.distributed_run:
                 reduced_loss = reduce_tensor(loss.data, n_gpus).item()
             else: