Repository: espnet/espnet
Issue: TTS BCE loss issue since labels is not modified exactly correctly in TTS Transformer

**Describe the bug**

In [epsnet2/tts/transformer](https://github.com/espnet/espnet/blob/4c7442ad8faad6a6d975337563fe7e1a795c00cb/espnet2/tts/transformer.py#L428) and [espnet/nets/pytorch_backend/e2e_tts_transformer.py](https://github.com/espnet/espnet/blob/4c7442ad8faad6a6d975337563fe7e1a795c00cb/espnet/nets/pytorch_backend/e2e_tts_transformer.py#L756)

```python
labels[:, -1] = 1.0  # make sure at least one frame has 1
```

It is not sure the stop frames when `self.reductor_factor > 1` and should be modified as

```python
import torch

labels = torch.scatter(labels, 1, (olens - 1).unsqueeze(1), 1.0)
```

**Example of the bug**

Given the case that 

```python
olens = torch.LongTensor([5, 3])
labels = torch.FloatTensor([
    [0, 0, 0, 0, 1],
    [0, 0, 1, 1, 1],
])
```

when `reduction_factor == 2`, follow the program, we would have

```python
olens = torch.LongTensor([4, 2])
labels = torch.FloatTensor([
    [0, 0, 0, 1], # correct
    [0, 0, 1, 1], # wrong, ``[0, 1, 1, 1]`` is correct
])
```

Here, we can see `olens` and `labels` are **mismatched**. So this may cause error while computing BCE loss.

Wish the report can be helpful.

Changed files:
espnet2/tts/transformer.py, espnet/nets/pytorch_backend/e2e_tts_transformer.py, espnet/nets/pytorch_backend/e2e_vc_tacotron2.py, test/espnet2/tts/test_transformer.py, espnet2/tts/tacotron2.py, espnet/nets/pytorch_backend/e2e_vc_transformer.py, test/espnet2/tts/test_tacotron2.py, espnet/nets/pytorch_backend/e2e_tts_tacotron2.py

diff --git a/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py b/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py
index c41dd4262..6954bab91 100644
--- a/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py
+++ b/espnet/nets/pytorch_backend/e2e_tts_tacotron2.py
@@ -736,11 +736,16 @@ class Tacotron2(TTSInterface, torch.nn.Module):
 
         # modifiy mod part of groundtruth
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_out = max(olens)
             ys = ys[:, :max_out]
             labels = labels[:, :max_out]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
 
         # caluculate taco2 loss
         l1_loss, mse_loss, bce_loss = self.taco2_loss(
diff --git a/espnet/nets/pytorch_backend/e2e_tts_transformer.py b/espnet/nets/pytorch_backend/e2e_tts_transformer.py
index 3628e7df0..7ed98e573 100644
--- a/espnet/nets/pytorch_backend/e2e_tts_transformer.py
+++ b/espnet/nets/pytorch_backend/e2e_tts_transformer.py
@@ -749,11 +749,16 @@ class Transformer(TTSInterface, torch.nn.Module):
 
         # modifiy mod part of groundtruth
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_olen = max(olens)
             ys = ys[:, :max_olen]
             labels = labels[:, :max_olen]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
 
         # caluculate loss values
         l1_loss, l2_loss, bce_loss = self.criterion(
diff --git a/espnet/nets/pytorch_backend/e2e_vc_tacotron2.py b/espnet/nets/pytorch_backend/e2e_vc_tacotron2.py
index e876c42cf..20ae50dea 100644
--- a/espnet/nets/pytorch_backend/e2e_vc_tacotron2.py
+++ b/espnet/nets/pytorch_backend/e2e_vc_tacotron2.py
@@ -572,11 +572,16 @@ class Tacotron2(TTSInterface, torch.nn.Module):
 
         # modifiy mod part of groundtruth
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_out = max(olens)
             ys = ys[:, :max_out]
             labels = labels[:, :max_out]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
         if self.encoder_reduction_factor > 1:
             ilens = ilens.new(
                 [ilen - ilen % self.encoder_reduction_factor for ilen in ilens]
diff --git a/espnet/nets/pytorch_backend/e2e_vc_transformer.py b/espnet/nets/pytorch_backend/e2e_vc_transformer.py
index 1b3c6913c..bcd40655c 100644
--- a/espnet/nets/pytorch_backend/e2e_vc_transformer.py
+++ b/espnet/nets/pytorch_backend/e2e_vc_transformer.py
@@ -718,11 +718,16 @@ class Transformer(TTSInterface, torch.nn.Module):
 
         # modifiy mod part of groundtruth
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_olen = max(olens)
             ys = ys[:, :max_olen]
             labels = labels[:, :max_olen]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
 
         # caluculate loss values
         l1_loss, l2_loss, bce_loss = self.criterion(
diff --git a/espnet2/tts/tacotron2.py b/espnet2/tts/tacotron2.py
index d5c8b3cc7..d94aaf180 100644
--- a/espnet2/tts/tacotron2.py
+++ b/espnet2/tts/tacotron2.py
@@ -305,11 +305,16 @@ class Tacotron2(AbsTTS):
 
         # modify mod part of groundtruth
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_out = max(olens)
             ys = ys[:, :max_out]
             labels = labels[:, :max_out]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
 
         # calculate taco2 loss
         l1_loss, mse_loss, bce_loss = self.taco2_loss(
diff --git a/espnet2/tts/transformer.py b/espnet2/tts/transformer.py
index f437c1845..94436c080 100644
--- a/espnet2/tts/transformer.py
+++ b/espnet2/tts/transformer.py
@@ -420,12 +420,17 @@ class Transformer(AbsTTS):
         # modifiy mod part of groundtruth
         olens_in = olens
         if self.reduction_factor > 1:
+            assert olens.ge(
+                self.reduction_factor
+            ).all(), "Output length must be greater than or equal to reduction factor."
             olens_in = olens.new([olen // self.reduction_factor for olen in olens])
             olens = olens.new([olen - olen % self.reduction_factor for olen in olens])
             max_olen = max(olens)
             ys = ys[:, :max_olen]
             labels = labels[:, :max_olen]
-            labels[:, -1] = 1.0  # make sure at least one frame has 1
+            labels = torch.scatter(
+                labels, 1, (olens - 1).unsqueeze(1), 1.0
+            )  # see #3388
 
         # caluculate loss values
         l1_loss, l2_loss, bce_loss = self.criterion(
diff --git a/test/espnet2/tts/test_tacotron2.py b/test/espnet2/tts/test_tacotron2.py
index 954a19896..272463579 100644
--- a/test/espnet2/tts/test_tacotron2.py
+++ b/test/espnet2/tts/test_tacotron2.py
@@ -58,8 +58,8 @@ def test_tacotron2(
     inputs = dict(
         text=torch.randint(0, 10, (2, 4)),
         text_lengths=torch.tensor([4, 1], dtype=torch.long),
-        speech=torch.randn(2, 3, 5),
-        speech_lengths=torch.tensor([3, 1], dtype=torch.long),
+        speech=torch.randn(2, 5, 5),
+        speech_lengths=torch.tensor([5, 3], dtype=torch.long),
     )
     if spk_embed_dim is not None:
         inputs.update(spembs=torch.randn(2, spk_embed_dim))
diff --git a/test/espnet2/tts/test_transformer.py b/test/espnet2/tts/test_transformer.py
index 68ba75f8e..43b58abd2 100644
--- a/test/espnet2/tts/test_transformer.py
+++ b/test/espnet2/tts/test_transformer.py
@@ -72,8 +72,8 @@ def test_tranformer(
     inputs = dict(
         text=torch.randint(0, 10, (2, 4)),
         text_lengths=torch.tensor([4, 1], dtype=torch.long),
-        speech=torch.randn(2, 3, 5),
-        speech_lengths=torch.tensor([3, 1], dtype=torch.long),
+        speech=torch.randn(2, 5, 5),
+        speech_lengths=torch.tensor([5, 3], dtype=torch.long),
     )
     if spk_embed_dim is not None:
         inputs.update(spembs=torch.randn(2, spk_embed_dim))