formatting, merge GST model with Tacotron

tests/test_layers.py

@@ -67,7 +67,8 @@ class DecoderTests(unittest.TestCase):
         assert output.shape[2] == 80 * 2, "size not {}".format(output.shape[2])
         assert stop_tokens.shape[0] == 4
 
-    def test_in_out_multispeaker(self):
+    @staticmethod
+    def test_in_out_multispeaker():
         layer = Decoder(
             in_features=256,
             memory_dim=80,

tests/test_tacotron_model.py

@@ -8,7 +8,6 @@ from torch import nn
 from TTS.utils.generic_utils import load_config
 from TTS.layers.losses import L1LossMasked
 from TTS.models.tacotron import Tacotron
-from TTS.models.tacotrongst import TacotronGST
 
 #pylint: disable=unused-variable
 
@@ -25,68 +24,72 @@ def count_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
 
-# class TacotronTrainTest(unittest.TestCase):
-    # def test_train_step(self):
-    #     input = 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['num_freq']).to(device)
-    #     mel_lengths = torch.randint(20, 30, (8, )).long().to(device)
-    #     stop_targets = torch.zeros(8, 30, 1).float().to(device)
-    #     speaker_ids = torch.randint(0, 5, (8, )).long().to(device)
+class TacotronTrainTest(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['num_freq']).to(device)
+        mel_lengths = torch.randint(20, 30, (8, )).long().to(device)
+        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
+        for idx in mel_lengths:
+            stop_targets[:, int(idx.item()):, 0] = 1.0
 
-    #     stop_targets = stop_targets.view(input.shape[0],
-    #                                      stop_targets.size(1) // c.r, -1)
-    #     stop_targets = (stop_targets.sum(2) >
-    #                     0.0).unsqueeze(2).float().squeeze()
+        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().to(device)
-    #     criterion_st = nn.BCEWithLogitsLoss().to(device)
-    #     model = Tacotron(
-    #         num_chars=32,
-    #         num_speakers=5, 
-    #         linear_dim=c.audio['num_freq'],
-    #         mel_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(" > 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, input_lengths, mel_spec, 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
-    #         # if count not in [145, 59]:
-    #         assert (param != param_ref).any(
-    #         ), "param {} with shape {} not updated!! \n{}\n{}".format(
-    #             count, param.shape, param, param_ref)
-    #         count += 1
+        criterion = L1LossMasked().to(device)
+        criterion_st = nn.BCEWithLogitsLoss().to(device)
+        model = Tacotron(
+            num_chars=32,
+            num_speakers=5,
+            linear_dim=c.audio['num_freq'],
+            mel_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(" > 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, 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
+            # if count not in [145, 59]:
+            assert (param != param_ref).any(
+            ), "param {} with shape {} not updated!! \n{}\n{}".format(
+                count, param.shape, param, param_ref)
+            count += 1
 
 
 class TacotronGSTTrainTest(unittest.TestCase):
-    def test_train_step(self):
-        input = torch.randint(0, 24, (8, 128)).long().to(device)
+    @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, 120, c.audio['num_mels']).to(device)
@@ -98,23 +101,26 @@ class TacotronGSTTrainTest(unittest.TestCase):
         for idx in mel_lengths:
             stop_targets[:, int(idx.item()):, 0] = 1.0
 
-        stop_targets = stop_targets.view(input.shape[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().to(device)
         criterion_st = nn.BCEWithLogitsLoss().to(device)
-        model = TacotronGST(
+        model = Tacotron(
             num_chars=32,
-            num_speakers=5, 
+            num_speakers=5,
+            gst=True,
             linear_dim=c.audio['num_freq'],
             mel_dim=c.audio['num_mels'],
             r=c.r,
-            memory_size=c.memory_size).to(device) #FIXME: missing num_speakers parameter to Tacotron ctor
+            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)))
+        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(),
@@ -124,7 +130,7 @@ class TacotronGSTTrainTest(unittest.TestCase):
         optimizer = optim.Adam(model.parameters(), lr=c.lr)
         for _ in range(10):
             mel_out, linear_out, align, stop_tokens = model.forward(
-                input, input_lengths, mel_spec, speaker_ids)
+                input_dummy, input_lengths, mel_spec, speaker_ids)
             optimizer.zero_grad()
             loss = criterion(mel_out, mel_spec, mel_lengths)
             stop_loss = criterion_st(stop_tokens, stop_targets)