refactor tacotron2 loss tests

yangarbiter · yangarbiter · commit 1996fec0cf58 · 2021-07-14T20:03:28.000Z
diff --git a/examples/pipeline_tacotron2/loss/loss_function.py b/examples/pipeline_tacotron2/loss/loss_function.py
@@ -38,6 +38,9 @@ class Tacotron2Loss(nn.Module):
     def __init__(self):
         super().__init__()
 
+        self.mse_loss = nn.MSELoss(reduction="mean")
+        self.bce_loss = nn.BCEWithLogitsLoss(reduction="mean")
+
     def forward(
         self,
         model_outputs: Tuple[Tensor, Tensor, Tensor],
@@ -56,27 +59,24 @@ def forward(
                     with shape (batch, mel, time).
                 (2) predicted mel spectrogram after the postnet (``mel_specgram_postnet``)
                     with shape (batch, mel, time), and
-                (3) the stop token prediction (``gate_out``) with shape (batch).
+                (3) the stop token prediction (``gate_out``) with shape (batch, ).
             targets (tuple of two Tensors): The ground truth mel spectrogram (batch, mel, time) and
-                stop token with shape (batch).
+                stop token with shape (batch, ).
 
         Returns:
-            mel_loss (Tensor): The mean MSE of the mel_specgram and ground truth mel spectrogram with shape (batch, ).
+            mel_loss (Tensor): The mean MSE of the mel_specgram and ground truth mel spectrogram
+                with shape ``torch.Size([])``.
             mel_postnet_loss (Tensor): The mean MSE of the mel_specgram_postnet and
-                ground truth mel spectrogram with shape (batch, ).
+                ground truth mel spectrogram with shape ``torch.Size([])``.
             gate_loss (Tensor): The mean binary cross entropy loss of
-                the prediction on the stop token with shape (batch, ).
+                the prediction on the stop token with shape ``torch.Size([])``.
         """
         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_specgram, mel_specgram_postnet, gate_out = model_outputs
         gate_out = gate_out.view(-1, 1)
-        mel_loss = nn.MSELoss(reduction="mean")(mel_specgram, mel_target)
-        mel_postnet_loss = nn.MSELoss(reduction="mean")(
-            mel_specgram_postnet, mel_target
-        )
-        gate_loss = nn.BCEWithLogitsLoss(reduction="mean")(gate_out, gate_target)
+        mel_loss = self.mse_loss(mel_specgram, mel_target)
+        mel_postnet_loss = self.mse_loss(mel_specgram_postnet, mel_target)
+        gate_loss = self.bce_loss(gate_out, gate_target)
         return mel_loss, mel_postnet_loss, gate_loss
diff --git a/examples/pipeline_tacotron2/loss/tacotron2_loss_cpu_test.py b/examples/pipeline_tacotron2/loss/tacotron2_loss_cpu_test.py
@@ -0,0 +1,23 @@
+import torch
+import unittest
+
+from .tacotron2_loss_impl import (
+    Tacotron2LossShapeTests,
+    Tacotron2LossTorchscriptTests,
+    Tacotron2LossGradcheckTests,
+)
+
+
+class TestTacotron2LossShapeFloat32CPU(unittest.TestCase, Tacotron2LossShapeTests):
+    dtype = torch.float32
+    device = torch.device("cpu")
+
+
+class TestTacotron2TorchsciptFloat32CPU(unittest.TestCase, Tacotron2LossTorchscriptTests):
+    dtype = torch.float32
+    device = torch.device("cpu")
+
+
+class TestTacotron2GradcheckFloat64CPU(unittest.TestCase, Tacotron2LossGradcheckTests):
+    dtype = torch.float64   # gradcheck needs a higher numerical accuracy
+    device = torch.device("cpu")
diff --git a/examples/pipeline_tacotron2/loss/tacotron2_loss_gpu_test.py b/examples/pipeline_tacotron2/loss/tacotron2_loss_gpu_test.py
@@ -0,0 +1,41 @@
+import os
+import unittest
+
+import torch
+
+from .tacotron2_loss_impl import (
+    Tacotron2LossShapeTests,
+    Tacotron2LossTorchscriptTests,
+    Tacotron2LossGradcheckTests,
+)
+
+
+def skipIfNoCuda(test_item):
+    if torch.cuda.is_available():
+        return test_item
+    force_cuda_test = os.environ.get("TORCHAUDIO_TEST_FORCE_CUDA", "0")
+    if force_cuda_test not in ["0", "1"]:
+        raise ValueError('"TORCHAUDIO_TEST_FORCE_CUDA" must be either "0" or "1".')
+    if force_cuda_test == "1":
+        raise RuntimeError(
+            '"TORCHAUDIO_TEST_FORCE_CUDA" is set but CUDA is not available.'
+        )
+    return unittest.skip("CUDA is not available.")(test_item)
+
+
+@skipIfNoCuda
+class TestTacotron2LossShapeFloat32CUDA(unittest.TestCase, Tacotron2LossShapeTests):
+    dtype = torch.float32
+    device = torch.device("cuda")
+
+
+@skipIfNoCuda
+class TestTacotron2TorchsciptFloat32CUDA(unittest.TestCase, Tacotron2LossTorchscriptTests):
+    dtype = torch.float32
+    device = torch.device("cuda")
+
+
+@skipIfNoCuda
+class TestTacotron2GradcheckFloat64CUDA(unittest.TestCase, Tacotron2LossGradcheckTests):
+    dtype = torch.float64   # gradcheck needs a higher numerical accuracy
+    device = torch.device("cuda")
diff --git a/examples/pipeline_tacotron2/loss/tacotron2_loss_impl.py b/examples/pipeline_tacotron2/loss/tacotron2_loss_impl.py
@@ -0,0 +1,143 @@
+import os
+import unittest
+import tempfile
+
+import torch
+from torch.autograd import gradcheck, gradgradcheck
+
+from .loss_function import Tacotron2Loss
+
+
+class TempDirMixin:
+    """Mixin to provide easy access to temp dir"""
+
+    temp_dir_ = None
+
+    @classmethod
+    def get_base_temp_dir(cls):
+        # If TORCHAUDIO_TEST_TEMP_DIR is set, use it instead of temporary directory.
+        # this is handy for debugging.
+        key = "TORCHAUDIO_TEST_TEMP_DIR"
+        if key in os.environ:
+            return os.environ[key]
+        if cls.temp_dir_ is None:
+            cls.temp_dir_ = tempfile.TemporaryDirectory()
+        return cls.temp_dir_.name
+
+    @classmethod
+    def tearDownClass(cls):
+        super().tearDownClass()
+        if cls.temp_dir_ is not None:
+            cls.temp_dir_.cleanup()
+            cls.temp_dir_ = None
+
+    def get_temp_path(self, *paths):
+        temp_dir = os.path.join(self.get_base_temp_dir(), self.id())
+        path = os.path.join(temp_dir, *paths)
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        return path
+
+
+class Tacotron2LossInputMixin(TempDirMixin):
+
+    def _get_inputs(self, n_mel=80, n_batch=16, max_mel_specgram_length=300):
+        mel_specgram = torch.rand(
+            n_batch, n_mel, max_mel_specgram_length, dtype=self.dtype, device=self.device
+        )
+        mel_specgram_postnet = torch.rand(
+            n_batch, n_mel, max_mel_specgram_length, dtype=self.dtype, device=self.device
+        )
+        gate_out = torch.rand(n_batch, dtype=self.dtype, device=self.device)
+        truth_mel_specgram = torch.rand(
+            n_batch, n_mel, max_mel_specgram_length, dtype=self.dtype, device=self.device
+        )
+        truth_gate_out = torch.rand(n_batch, dtype=self.dtype, device=self.device)
+
+        truth_mel_specgram.requires_grad = False
+        truth_gate_out.requires_grad = False
+
+        return (
+            mel_specgram,
+            mel_specgram_postnet,
+            gate_out,
+            truth_mel_specgram,
+            truth_gate_out,
+        )
+
+
+class Tacotron2LossShapeTests(Tacotron2LossInputMixin):
+
+    def test_tacotron2_loss_shape(self):
+        f"""Validate the output shape of Tacotron2Loss."""
+        n_batch = 16
+
+        (
+            mel_specgram,
+            mel_specgram_postnet,
+            gate_out,
+            truth_mel_specgram,
+            truth_gate_out,
+        ) = self._get_inputs(n_batch=n_batch)
+
+        mel_loss, mel_postnet_loss, gate_loss  = Tacotron2Loss()(
+            (mel_specgram, mel_specgram_postnet, gate_out),
+            (truth_mel_specgram, truth_gate_out)
+        )
+
+        self.assertEqual(mel_loss.size(), torch.Size([]))
+        self.assertEqual(mel_postnet_loss.size(), torch.Size([]))
+        self.assertEqual(gate_loss.size(), torch.Size([]))
+
+
+class Tacotron2LossTorchscriptTests(Tacotron2LossInputMixin):
+
+    def _assert_torchscript_consistency(self, fn, tensors):
+        path = self.get_temp_path("func.zip")
+        torch.jit.script(fn).save(path)
+        ts_func = torch.jit.load(path)
+
+        output = fn(tensors[:3], tensors[3:])
+        ts_output = ts_func(tensors[:3], tensors[3:])
+
+        self.assertEqual(ts_output, output)
+
+    def test_tacotron2_loss_torchscript_consistency(self):
+        f"""Validate the torchscript consistency of Tacotron2Loss."""
+
+        loss_fn = Tacotron2Loss()
+        self._assert_torchscript_consistency(loss_fn, self._get_inputs())
+
+
+class Tacotron2LossGradcheckTests(Tacotron2LossInputMixin):
+
+    def test_tacotron2_loss_gradcheck(self):
+        f"""Performing gradient check on Tacotron2Loss."""
+        (
+            mel_specgram,
+            mel_specgram_postnet,
+            gate_out,
+            truth_mel_specgram,
+            truth_gate_out,
+        ) = self._get_inputs()
+
+        mel_specgram.requires_grad_(True)
+        mel_specgram_postnet.requires_grad_(True)
+        gate_out.requires_grad_(True)
+
+        def _fn(mel_specgram, mel_specgram_postnet, gate_out, truth_mel_specgram, truth_gate_out):
+            loss_fn = Tacotron2Loss()
+            return loss_fn(
+                (mel_specgram, mel_specgram_postnet, gate_out),
+                (truth_mel_specgram, truth_gate_out),
+            )
+
+        gradcheck(
+            _fn,
+            (mel_specgram, mel_specgram_postnet, gate_out, truth_mel_specgram, truth_gate_out),
+            fast_mode=True,
+        )
+        gradgradcheck(
+            _fn,
+            (mel_specgram, mel_specgram_postnet, gate_out, truth_mel_specgram, truth_gate_out),
+            fast_mode=True,
+        )
diff --git a/examples/pipeline_tacotron2/loss/test_tacotron2_loss.py b/examples/pipeline_tacotron2/loss/test_tacotron2_loss.py
