travis unit tests fix and add Tacotron and Tacotron 2 GST and MultiSpeaker Tests

tests/inputs/test_config.json

@@ -53,6 +53,7 @@
     "max_seq_len": 300,
     "log_dir": "tests/outputs/",
 
+<<<<<<< HEAD
     "use_speaker_embedding": false,     
     "use_gst": false,       			    
     "gst":	{			                
@@ -61,4 +62,18 @@
         "gst_num_heads": 4,
         "gst_style_tokens": 10
 	}
+=======
+    // MULTI-SPEAKER and GST
+    "use_speaker_embedding": false,     // use speaker embedding to enable multi-speaker learning.
+    "use_gst": true,       			    // use global style tokens
+    "gst":	{			                // gst parameter if gst is enabled
+        "gst_style_input": null,        // Condition the style input either on a 
+                                        // -> 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_embedding_dim": 512,       
+        "gst_num_heads": 4,
+        "gst_style_tokens": 10
+        }  
+>>>>>>> travis unit tests fix and add Tacotron and Tacotron 2 GST and MultiSpeaker Tests
     }

tests/test_layers.py

@@ -58,8 +58,7 @@ class DecoderTests(unittest.TestCase):
             trans_agent=True,
             forward_attn_mask=True,
             location_attn=True,
-            separate_stopnet=True,
-            speaker_embedding_dim=0)
+            separate_stopnet=True)
         dummy_input = T.rand(4, 8, 256)
         dummy_memory = T.rand(4, 2, 80)
 
@@ -71,38 +70,6 @@ class DecoderTests(unittest.TestCase):
         assert output.shape[2] == 2, "size not {}".format(output.shape[2])
         assert stop_tokens.shape[0] == 4
 
-    @staticmethod
-    def test_in_out_multispeaker():
-        layer = Decoder(
-            in_channels=256,
-            frame_channels=80,
-            r=2,
-            memory_size=4,
-            attn_windowing=False,
-            attn_norm="sigmoid",
-            attn_K=5,
-            attn_type="graves",
-            prenet_type='original',
-            prenet_dropout=True,
-            forward_attn=True,
-            trans_agent=True,
-            forward_attn_mask=True,
-            location_attn=True,
-            separate_stopnet=True,
-            speaker_embedding_dim=80)
-        dummy_input = T.rand(4, 8, 256)
-        dummy_memory = T.rand(4, 2, 80)
-        dummy_embed = T.rand(4, 80)
-
-        output, alignment, stop_tokens = layer(
-            dummy_input, dummy_memory, mask=None, speaker_embeddings=dummy_embed)
-
-        assert output.shape[0] == 4
-        assert output.shape[1] == 80, "size not {}".format(output.shape[1])
-        assert output.shape[2] == 2, "size not {}".format(output.shape[2])
-        assert stop_tokens.shape[0] == 4
-
-
 class EncoderTests(unittest.TestCase):
     def test_in_out(self):  #pylint: disable=no-self-use
         layer = Encoder(128)

tests/test_tacotron2_model.py

@@ -9,6 +9,7 @@ from torch import nn, optim
 from mozilla_voice_tts.tts.layers.losses import MSELossMasked
 from mozilla_voice_tts.tts.models.tacotron2 import Tacotron2
 from mozilla_voice_tts.utils.io import load_config
+from mozilla_voice_tts.utils.audio import AudioProcessor
 
 #pylint: disable=unused-variable
 
@@ -18,14 +19,12 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 c = load_config(os.path.join(get_tests_input_path(), 'test_config.json'))
 
+ap = AudioProcessor(**c.audio)
+WAV_FILE = os.path.join(get_tests_input_path(), "example_1.wav")
+
 
 class TacotronTrainTest(unittest.TestCase):
-<<<<<<< HEAD
     def test_train_step(self):  # pylint: disable=no-self-use
-=======
-    @staticmethod
-    def test_train_step():
->>>>>>> small gst config change
         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]
@@ -75,3 +74,113 @@ class TacotronTrainTest(unittest.TestCase):
             ), "param {} with shape {} not updated!! \n{}\n{}".format(
                 count, param.shape, param, param_ref)
             count += 1
+
+class TacotronGSTTrainTest(unittest.TestCase):
+    @staticmethod
+    def test_train_step():
+        # with random gst mel style
+        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_ids = torch.randint(0, 5, (8, )).long().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, 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']).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(10):
+            mel_out, mel_postnet_out, align, stop_tokens = model.forward(
+                input_dummy, input_lengths, mel_spec, mel_lengths, speaker_ids)
+            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':
+                #print(param.grad)
+                continue
+            assert (param != param_ref).any(
+            ), "param {} {} with shape {} not updated!! \n{}\n{}".format(
+                name, count, param.shape, param, param_ref)
+            count += 1
+
+        # with file gst style
+        mel_spec = torch.FloatTensor(ap.melspectrogram(ap.load_wav(WAV_FILE)))[:, :30].unsqueeze(0).transpose(1, 2).to(device)
+        mel_spec = mel_spec.repeat(8, 1, 1)
+        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_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_ids = torch.randint(0, 5, (8, )).long().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, 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']).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(10):
+            mel_out, mel_postnet_out, align, stop_tokens = model.forward(
+                input_dummy, input_lengths, mel_spec, mel_lengths, speaker_ids)
+            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':
+                #print(param.grad)
+                continue
+            assert (param != param_ref).any(
+            ), "param {} {} with shape {} not updated!! \n{}\n{}".format(
+                name, count, param.shape, param, param_ref)
+            count += 1

tests/test_tacotron_model.py

@@ -9,6 +9,7 @@ from torch import nn, optim
 from mozilla_voice_tts.tts.layers.losses import L1LossMasked
 from mozilla_voice_tts.tts.models.tacotron import Tacotron
 from mozilla_voice_tts.utils.io import load_config
+from mozilla_voice_tts.utils.audio import AudioProcessor
 
 #pylint: disable=unused-variable
 
@@ -18,6 +19,9 @@ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 c = load_config(os.path.join(get_tests_input_path(), 'test_config.json'))
 
+ap = AudioProcessor(**c.audio)
+WAV_FILE = os.path.join(get_tests_input_path(), "example_1.wav")
+
 
 def count_parameters(model):
     r"""Count number of trainable parameters in a network"""
@@ -85,10 +89,10 @@ class TacotronTrainTest(unittest.TestCase):
                 count, param.shape, param, param_ref)
             count += 1
 
-
 class TacotronGSTTrainTest(unittest.TestCase):
     @staticmethod
     def test_train_step():
+        # with random gst mel style
         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
@@ -113,13 +117,82 @@ class TacotronGSTTrainTest(unittest.TestCase):
             num_chars=32,
             num_speakers=5,
             gst=True,
-            postnet_output_dim=c.audio['num_freq'],
+            gst_embedding_dim=c.gst['gst_embedding_dim'],
+            gst_num_heads=c.gst['gst_num_heads'],
+            gst_style_tokens=c.gst['gst_style_tokens'],
+            postnet_output_dim=c.audio['fft_size'],
             decoder_output_dim=c.audio['num_mels'],
             r=c.r,
             memory_size=c.memory_size
         ).to(device)  #FIXME: missing num_speakers parameter to Tacotron ctor
         model.train()
-        print(model)
+        # print(model)
+        print(" > Num parameters for Tacotron GST 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(10):
+            mel_out, linear_out, align, stop_tokens = model.forward(
+                input_dummy, input_lengths, mel_spec, mel_lengths, speaker_ids)
+            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 param, param_ref in zip(model.parameters(),
+                                    model_ref.parameters()):
+            # ignore pre-higway layer since it works conditional
+            assert (param != param_ref).any(
+            ), "param {} with shape {} not updated!! \n{}\n{}".format(
+                count, param.shape, param, param_ref)
+            count += 1
+
+        # with file gst style
+        mel_spec = torch.FloatTensor(ap.melspectrogram(ap.load_wav(WAV_FILE)))[:, :120].unsqueeze(0).transpose(1, 2).to(device)
+        mel_spec = mel_spec.repeat(8, 1, 1)
+
+        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
+        linear_spec = torch.rand(8, mel_spec.size(1), c.audio['fft_size']).to(device)
+        mel_lengths = torch.randint(20, mel_spec.size(1), (8, )).long().to(device)
+        mel_lengths[-1] = mel_spec.size(1)
+        stop_targets = torch.zeros(8, mel_spec.size(1), 1).float().to(device)
+        speaker_ids = torch.randint(0, 5, (8, )).long().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,
+            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'],
+            postnet_output_dim=c.audio['fft_size'],
+            decoder_output_dim=c.audio['num_mels'],
+            r=c.r,
+            memory_size=c.memory_size
+        ).to(device)  #FIXME: missing num_speakers parameter to Tacotron ctor
+        model.train()
+        # print(model)
         print(" > Num parameters for Tacotron GST model:%s" %
               (count_parameters(model)))
         model_ref = copy.deepcopy(model)