add unit tests for SC-GST

tests/inputs/test_config.json

@@ -61,6 +61,7 @@
                                         // -> wave file [path to wave] or
                                         // -> dictionary using the style tokens {'token1': 'value', 'token2': 'value'} example {"0": 0.15, "1": 0.15, "5": -0.15}
                                         // with the dictionary being len(dict) <= len(gst_style_tokens).
+        "gst_use_speaker_embedding": true, // if true pass speaker embedding in attention input GST.
         "gst_embedding_dim": 512,
         "gst_num_heads": 4,
         "gst_style_tokens": 10

tests/inputs/test_train_config.json

@@ -140,6 +140,7 @@
                                         // -> wave file [path to wave] or
                                         // -> dictionary using the style tokens {'token1': 'value', 'token2': 'value'} example {"0": 0.15, "1": 0.15, "5": -0.15}
                                         // with the dictionary being len(dict) == len(gst_style_tokens).
+        "gst_use_speaker_embedding": true, // if true pass speaker embedding in attention input GST.
         "gst_embedding_dim": 512,
         "gst_num_heads": 4,
         "gst_style_tokens": 10

tests/outputs/dummy_model_config.json

@@ -93,6 +93,7 @@
                                         // -> wave file [path to wave] or
                                         // -> dictionary using the style tokens {'token1': 'value', 'token2': 'value'} example {"0": 0.15, "1": 0.15, "5": -0.15}
                                         // with the dictionary being len(dict) <= len(gst_style_tokens).
+        "gst_use_speaker_embedding": true, // if true pass speaker embedding in attention input GST.
         "gst_embedding_dim": 512,
         "gst_num_heads": 4,
         "gst_style_tokens": 10

tests/test_tacotron2_model.py

@@ -238,3 +238,58 @@ class TacotronGSTTrainTest(unittest.TestCase):
             ), "param {} {} with shape {} not updated!! \n{}\n{}".format(
                 name, count, param.shape, param, param_ref)
             count += 1
+
+class SCGSTMultiSpeakeTacotronTrainTest(unittest.TestCase):
+    @staticmethod
+    def test_train_step():
+        input_dummy = torch.randint(0, 24, (8, 128)).long().to(device)
+        input_lengths = torch.randint(100, 128, (8, )).long().to(device)
+        input_lengths = torch.sort(input_lengths, descending=True)[0]
+        mel_spec = torch.rand(8, 30, c.audio['num_mels']).to(device)
+        mel_postnet_spec = torch.rand(8, 30, c.audio['num_mels']).to(device)
+        mel_lengths = torch.randint(20, 30, (8, )).long().to(device)
+        mel_lengths[0] = 30
+        stop_targets = torch.zeros(8, 30, 1).float().to(device)
+        speaker_embeddings = torch.rand(8, 55).to(device)
+
+        for idx in mel_lengths:
+            stop_targets[:, int(idx.item()):, 0] = 1.0
+
+        stop_targets = stop_targets.view(input_dummy.shape[0],
+                                         stop_targets.size(1) // c.r, -1)
+        stop_targets = (stop_targets.sum(2) > 0.0).unsqueeze(2).float().squeeze()
+        criterion = MSELossMasked(seq_len_norm=False).to(device)
+        criterion_st = nn.BCEWithLogitsLoss().to(device)
+        model = Tacotron2(num_chars=24, r=c.r, num_speakers=5, speaker_embedding_dim=55, gst=True, gst_embedding_dim=c.gst['gst_embedding_dim'], gst_num_heads=c.gst['gst_num_heads'], gst_style_tokens=c.gst['gst_style_tokens'], gst_use_speaker_embedding=c.gst['gst_use_speaker_embedding']).to(device)
+        model.train()
+        model_ref = copy.deepcopy(model)
+        count = 0
+        for param, param_ref in zip(model.parameters(),
+                                    model_ref.parameters()):
+            assert (param - param_ref).sum() == 0, param
+            count += 1
+        optimizer = optim.Adam(model.parameters(), lr=c.lr)
+        for i in range(5):
+            mel_out, mel_postnet_out, align, stop_tokens = model.forward(
+                input_dummy, input_lengths, mel_spec, mel_lengths, speaker_embeddings=speaker_embeddings)
+            assert torch.sigmoid(stop_tokens).data.max() <= 1.0
+            assert torch.sigmoid(stop_tokens).data.min() >= 0.0
+            optimizer.zero_grad()
+            loss = criterion(mel_out, mel_spec, mel_lengths)
+            stop_loss = criterion_st(stop_tokens, stop_targets)
+            loss = loss + criterion(mel_postnet_out, mel_postnet_spec, mel_lengths) + stop_loss
+            loss.backward()
+            optimizer.step()
+        # check parameter changes
+        count = 0
+        for name_param, param_ref in zip(model.named_parameters(),
+                                         model_ref.parameters()):
+            # ignore pre-higway layer since it works conditional
+            # if count not in [145, 59]:
+            name, param = name_param
+            if name == 'gst_layer.encoder.recurrence.weight_hh_l0':
+                continue
+            assert (param != param_ref).any(
+            ), "param {} with shape {} not updated!! \n{}\n{}".format(
+                count, param.shape, param, param_ref)
+            count += 1

tests/test_tacotron_model.py

@@ -284,3 +284,75 @@ class TacotronGSTTrainTest(unittest.TestCase):
             ), "param {} with shape {} not updated!! \n{}\n{}".format(
                 count, param.shape, param, param_ref)
             count += 1
+
+class SCGSTMultiSpeakeTacotronTrainTest(unittest.TestCase):
+    @staticmethod
+    def test_train_step():
+        input_dummy = torch.randint(0, 24, (8, 128)).long().to(device)
+        input_lengths = torch.randint(100, 129, (8, )).long().to(device)
+        input_lengths[-1] = 128
+        mel_spec = torch.rand(8, 30, c.audio['num_mels']).to(device)
+        linear_spec = torch.rand(8, 30, c.audio['fft_size']).to(device)
+        mel_lengths = torch.randint(20, 30, (8, )).long().to(device)
+        stop_targets = torch.zeros(8, 30, 1).float().to(device)
+        speaker_embeddings = torch.rand(8, 55).to(device)
+
+        for idx in mel_lengths:
+            stop_targets[:, int(idx.item()):, 0] = 1.0
+
+        stop_targets = stop_targets.view(input_dummy.shape[0],
+                                         stop_targets.size(1) // c.r, -1)
+        stop_targets = (stop_targets.sum(2) >
+                        0.0).unsqueeze(2).float().squeeze()
+
+        criterion = L1LossMasked(seq_len_norm=False).to(device)
+        criterion_st = nn.BCEWithLogitsLoss().to(device)
+        model = Tacotron(
+            num_chars=32,
+            num_speakers=5,
+            postnet_output_dim=c.audio['fft_size'],
+            decoder_output_dim=c.audio['num_mels'],
+            gst=True,
+            gst_embedding_dim=c.gst['gst_embedding_dim'],
+            gst_num_heads=c.gst['gst_num_heads'],
+            gst_style_tokens=c.gst['gst_style_tokens'],
+            gst_use_speaker_embedding=c.gst['gst_use_speaker_embedding'],
+            r=c.r,
+            memory_size=c.memory_size,
+            speaker_embedding_dim=55,
+        ).to(device)  #FIXME: missing num_speakers parameter to Tacotron ctor
+        model.train()
+        print(" > Num parameters for Tacotron model:%s" %
+              (count_parameters(model)))
+        model_ref = copy.deepcopy(model)
+        count = 0
+        for param, param_ref in zip(model.parameters(),
+                                    model_ref.parameters()):
+            assert (param - param_ref).sum() == 0, param
+            count += 1
+        optimizer = optim.Adam(model.parameters(), lr=c.lr)
+        for _ in range(5):
+            mel_out, linear_out, align, stop_tokens = model.forward(
+                input_dummy, input_lengths, mel_spec, mel_lengths,
+                speaker_embeddings=speaker_embeddings)
+            optimizer.zero_grad()
+            loss = criterion(mel_out, mel_spec, mel_lengths)
+            stop_loss = criterion_st(stop_tokens, stop_targets)
+            loss = loss + criterion(linear_out, linear_spec,
+                                    mel_lengths) + stop_loss
+            loss.backward()
+            optimizer.step()
+        # check parameter changes
+        count = 0
+        for name_param, param_ref in zip(model.named_parameters(),
+                                         model_ref.parameters()):
+            # ignore pre-higway layer since it works conditional
+            # if count not in [145, 59]:
+            name, param = name_param
+            if name == 'gst_layer.encoder.recurrence.weight_hh_l0':
+                continue
+            assert (param != param_ref).any(
+            ), "param {} with shape {} not updated!! \n{}\n{}".format(
+                count, param.shape, param, param_ref)
+            count += 1
+