add unit tests for SoftmaxAngleProtoLoss and ResnetSpeakerEncoder and bugfix

tests/test_speaker_encoder.py

@@ -3,13 +3,13 @@ import unittest
 import torch as T
 
 from tests import get_tests_input_path
-from TTS.speaker_encoder.losses import AngleProtoLoss, GE2ELoss
+from TTS.speaker_encoder.losses import AngleProtoLoss, GE2ELoss, SoftmaxAngleProtoLoss
 from TTS.speaker_encoder.models.lstm import LSTMSpeakerEncoder
-# from TTS.speaker_encoder.models.resnet import ResNetSpeakerEncoder
+from TTS.speaker_encoder.models.resnet import ResNetSpeakerEncoder
 file_path = get_tests_input_path()
 
 
-class SpeakerEncoderTests(unittest.TestCase):
+class LSTMSpeakerEncoderTests(unittest.TestCase):
     # pylint: disable=R0201
     def test_in_out(self):
         dummy_input = T.rand(4, 20, 80)  # B x T x D
@@ -39,6 +39,31 @@ class SpeakerEncoderTests(unittest.TestCase):
         assert output.shape[1] == 256
         assert len(output.shape) == 2
 
+class ResNetSpeakerEncoderTests(unittest.TestCase):
+    # pylint: disable=R0201
+    def test_in_out(self):
+        dummy_input = T.rand(4, 20, 80)  # B x T x D
+        dummy_hidden = [T.rand(2, 4, 128), T.rand(2, 4, 128)]
+        model = ResNetSpeakerEncoder(input_dim=80, proj_dim=256)
+        # computing d vectors
+        output = model.forward(dummy_input)
+        assert output.shape[0] == 4
+        assert output.shape[1] == 256
+        output = model.forward(dummy_input, l2_norm=True)
+        assert output.shape[0] == 4
+        assert output.shape[1] == 256
+
+        # check normalization
+        output_norm = T.nn.functional.normalize(output, dim=1, p=2)
+        assert_diff = (output_norm - output).sum().item()
+        assert output.type() == "torch.FloatTensor"
+        assert abs(assert_diff) < 1e-4, f" [!] output_norm has wrong values - {assert_diff}"
+        # compute d for a given batch
+        dummy_input = T.rand(1, 240, 80)  # B x T x D
+        output = model.compute_embedding(dummy_input, num_frames=160, num_eval=10)
+        assert output.shape[0] == 1
+        assert output.shape[1] == 256
+        assert len(output.shape) == 2
 
 class GE2ELossTests(unittest.TestCase):
     # pylint: disable=R0201
@@ -67,7 +92,6 @@ class GE2ELossTests(unittest.TestCase):
         output = loss.forward(dummy_input)
         assert output.item() < 0.005
 
-
 class AngleProtoLossTests(unittest.TestCase):
     # pylint: disable=R0201
     def test_in_out(self):
@@ -96,3 +120,24 @@ class AngleProtoLossTests(unittest.TestCase):
         loss = AngleProtoLoss()
         output = loss.forward(dummy_input)
         assert output.item() < 0.005
+
+class SoftmaxAngleProtoLossTests(unittest.TestCase):
+    # pylint: disable=R0201
+    def test_in_out(self):
+
+        embedding_dim = 64
+        num_speakers = 5
+        batch_size = 4
+
+        dummy_label = T.randint(low=0, high=num_speakers, size=(batch_size, num_speakers))
+        # check random input
+        dummy_input = T.rand(batch_size, num_speakers, embedding_dim)  # num_speaker x num_utterance x dim
+        loss = SoftmaxAngleProtoLoss(embedding_dim=embedding_dim, n_speakers=num_speakers)
+        output = loss.forward(dummy_input, dummy_label)
+        assert output.item() >= 0.0
+
+        # check all zeros
+        dummy_input = T.ones(batch_size, num_speakers, embedding_dim)  # num_speaker x num_utterance x dim
+        loss = SoftmaxAngleProtoLoss(embedding_dim=embedding_dim, n_speakers=num_speakers)
+        output = loss.forward(dummy_input, dummy_label)
+        assert output.item() >= 0.0