mirror of https://github.com/coqui-ai/TTS.git
Implement ForwardTTSE2E tests
This commit is contained in:
Eren Gölge
parent
ade84aa124
commit
2c0cd0ddd5
|
|
@ -0,0 +1,366 @@
|
|||
import copy
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from trainer.logging.tensorboard_logger import TensorboardLogger
|
||||
|
||||
from tests import assertHasAttr, assertHasNotAttr, get_tests_data_path, get_tests_input_path, get_tests_output_path
|
||||
from TTS.tts.configs.fast_pitch_e2e_config import FastPitchE2EConfig
|
||||
from TTS.tts.models.forward_tts_e2e import ForwardTTSE2E, ForwardTTSE2EArgs
|
||||
|
||||
LANG_FILE = os.path.join(get_tests_input_path(), "language_ids.json")
|
||||
SPEAKER_ENCODER_CONFIG = os.path.join(get_tests_input_path(), "test_speaker_encoder_config.json")
|
||||
WAV_FILE = os.path.join(get_tests_input_path(), "example_1.wav")
|
||||
|
||||
|
||||
torch.manual_seed(1)
|
||||
use_cuda = torch.cuda.is_available()
|
||||
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
# pylint: disable=no-self-use
|
||||
class TestFastPitchE2E(unittest.TestCase):
|
||||
def _create_inputs(self, config, batch_size=2):
|
||||
input_dummy = torch.randint(0, 24, (batch_size, 128)).long().to(device)
|
||||
input_lengths = torch.randint(100, 129, (batch_size,)).long().to(device)
|
||||
input_lengths[-1] = 128
|
||||
spec = torch.rand(batch_size, 30, config.audio["num_mels"]).to(device)
|
||||
# spec = torch.rand(batch_size, config.audio["num_mels"], 30).to(device)
|
||||
spec_lengths = torch.randint(20, 30, (batch_size,)).long().to(device)
|
||||
spec_lengths[-1] = spec.size(1)
|
||||
waveform = torch.rand(batch_size, 1, spec.size(1) * config.audio["hop_length"]).to(device)
|
||||
pitch = torch.rand(batch_size, 1, spec.size(1)).to(device)
|
||||
return input_dummy, input_lengths, spec, spec_lengths, waveform, pitch
|
||||
|
||||
def _check_forward_outputs(self, config, output_dict, batch_size=2):
|
||||
self.assertEqual(
|
||||
output_dict["model_outputs"].shape[2], config.model_args.spec_segment_size * config.audio["hop_length"]
|
||||
)
|
||||
self.assertEqual(output_dict["alignments"].shape, (batch_size, 30, 128))
|
||||
self.assertEqual(output_dict["alignments"].max(), 1)
|
||||
self.assertEqual(output_dict["alignments"].min(), 0)
|
||||
self.assertEqual(
|
||||
output_dict["waveform_seg"].shape[2], config.model_args.spec_segment_size * config.audio["hop_length"]
|
||||
)
|
||||
|
||||
def _check_inference_outputs(self, outputs, input_dummy, batch_size=1):
|
||||
feat_len = outputs["encoder_outputs"].shape[1]
|
||||
self.assertEqual(outputs["model_outputs"].shape[:2], (batch_size, 1)) # we don't know the channel dimension
|
||||
self.assertEqual(outputs["alignments"].shape, (batch_size, input_dummy.shape[1], feat_len))
|
||||
|
||||
@staticmethod
|
||||
def _check_parameter_changes(model, model_ref):
|
||||
count = 0
|
||||
for item1, item2 in zip(model.named_parameters(), model_ref.named_parameters()):
|
||||
name = item1[0]
|
||||
param = item1[1]
|
||||
param_ref = item2[1]
|
||||
assert (param != param_ref).any(), "param {} with shape {} not updated!! \n{}\n{}".format(
|
||||
name, param.shape, param, param_ref
|
||||
)
|
||||
count = count + 1
|
||||
|
||||
def _create_batch(self, config, batch_size):
|
||||
input_dummy, input_lengths, spec, spec_lengths, waveform, pitch = self._create_inputs(config, batch_size)
|
||||
batch = {}
|
||||
batch["text_input"] = input_dummy
|
||||
batch["text_lengths"] = input_lengths
|
||||
batch["mel_lengths"] = spec_lengths
|
||||
batch["mel_input"] = spec
|
||||
batch["waveform"] = waveform.transpose(1, 2) # B x C X T -> B x T x C
|
||||
batch["d_vectors"] = None
|
||||
batch["speaker_ids"] = None
|
||||
batch["language_ids"] = None
|
||||
batch["pitch"] = pitch
|
||||
return batch
|
||||
|
||||
# def test_init_multispeaker(self):
|
||||
|
||||
# num_speakers = 10
|
||||
# model_args = ForwardTTSE2EArgs()
|
||||
# model_args.num_speakers = num_speakers
|
||||
# model_args.use_speaker_embedding = True
|
||||
# model = ForwardTTSE2E(model_args)
|
||||
# assertHasAttr(self, model.encoder_model, "emb_g")
|
||||
|
||||
# model_args = ForwardTTSE2EArgs()
|
||||
# model_args.num_speakers = 0
|
||||
# model_args.use_speaker_embedding = True
|
||||
# model = ForwardTTSE2E(model_args)
|
||||
# assertHasNotAttr(self, model.encoder_model, "emb_g")
|
||||
|
||||
# model_args = ForwardTTSE2EArgs()
|
||||
# model_args.num_speakers = 10
|
||||
# model_args.use_speaker_embedding = False
|
||||
# model = ForwardTTSE2E(model_args)
|
||||
# assertHasNotAttr(self, model.encoder_model, "emb_g")
|
||||
|
||||
# model_args = ForwardTTSE2EArgs(d_vector_dim=101, use_d_vector_file=True)
|
||||
# model = ForwardTTSE2E(model_args)
|
||||
# self.assertEqual(model.encoder_model.embedded_speaker_dim, 101)
|
||||
|
||||
# def test_init_multilingual(self):
|
||||
# """TODO"""
|
||||
|
||||
# def test_get_aux_input(self):
|
||||
# aux_input = {"speaker_ids": None, "style_wav": None, "d_vectors": None, "language_ids": None}
|
||||
# model_args = ForwardTTSE2EArgs()
|
||||
# model = ForwardTTSE2E(model_args)
|
||||
# aux_out = model.get_aux_input(aux_input)
|
||||
|
||||
# speaker_id = torch.randint(10, (1,))
|
||||
# language_id = torch.randint(10, (1,))
|
||||
# d_vector = torch.rand(1, 128)
|
||||
# aux_input = {"speaker_ids": speaker_id, "style_wav": None, "d_vectors": d_vector, "language_ids": language_id}
|
||||
# aux_out = model.get_aux_input(aux_input)
|
||||
# self.assertEqual(aux_out["speaker_ids"].shape, speaker_id.shape)
|
||||
# self.assertEqual(aux_out["language_ids"].shape, language_id.shape)
|
||||
# self.assertEqual(aux_out["d_vectors"].shape, d_vector.unsqueeze(0).transpose(2, 1).shape)
|
||||
|
||||
# def test_forward(self):
|
||||
# model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# input_dummy, input_lengths, spec, spec_lengths, waveform, pitch = self._create_inputs(config)
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
# output_dict = model.forward(
|
||||
# x=input_dummy, x_lengths=input_lengths, spec=spec, spec_lengths=spec_lengths, waveform=waveform, pitch=pitch
|
||||
# )
|
||||
# self._check_forward_outputs(config, output_dict)
|
||||
|
||||
# def test_multispeaker_forward(self):
|
||||
# batch_size = 2
|
||||
# num_speakers = 10
|
||||
# model_args = ForwardTTSE2EArgs(
|
||||
# spec_segment_size=10, num_speakers=num_speakers, use_speaker_embedding=True
|
||||
# )
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# config.model_args.spec_segment_size = 10
|
||||
|
||||
# input_dummy, input_lengths, spec, spec_lengths, waveform, pitch = self._create_inputs(
|
||||
# config, batch_size=batch_size
|
||||
# )
|
||||
# speaker_ids = torch.randint(0, num_speakers, (batch_size,)).long().to(device)
|
||||
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
# output_dict = model.forward(
|
||||
# x=input_dummy,
|
||||
# x_lengths=input_lengths,
|
||||
# spec=spec,
|
||||
# spec_lengths=spec_lengths,
|
||||
# waveform=waveform,
|
||||
# pitch=pitch,
|
||||
# aux_input={"speaker_ids": speaker_ids},
|
||||
# )
|
||||
# self._check_forward_outputs(config, output_dict)
|
||||
|
||||
# def test_d_vector_forward(self):
|
||||
# batch_size = 2
|
||||
# model_args = ForwardTTSE2EArgs(
|
||||
# spec_segment_size=10, use_d_vector_file=True, d_vector_dim=256
|
||||
# )
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# config.model_args.spec_segment_size = 10
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
# model.train()
|
||||
# input_dummy, input_lengths, spec, spec_lengths, waveform, pitch = self._create_inputs(
|
||||
# config, batch_size=batch_size
|
||||
# )
|
||||
# d_vectors = torch.randn(batch_size, 256).to(device)
|
||||
# output_dict = model.forward(
|
||||
# x=input_dummy,
|
||||
# x_lengths=input_lengths,
|
||||
# spec=spec,
|
||||
# spec_lengths=spec_lengths,
|
||||
# waveform=waveform,
|
||||
# pitch=pitch,
|
||||
# aux_input={"d_vectors": d_vectors},
|
||||
# )
|
||||
# self._check_forward_outputs(config, output_dict)
|
||||
|
||||
# # def test_multilingual_forward(self):
|
||||
# # """TODO"""
|
||||
|
||||
# def test_inference(self):
|
||||
# model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
# model.eval()
|
||||
|
||||
# batch_size = 1
|
||||
# input_dummy, *_ = self._create_inputs(config, batch_size=batch_size)
|
||||
# outputs = model.inference(input_dummy.to(device))
|
||||
# self._check_inference_outputs(outputs, input_dummy, batch_size=batch_size)
|
||||
|
||||
# # TODO implemented batched inferenece
|
||||
# # batch_size = 2
|
||||
# # input_dummy, input_lengths, *_ = self._create_inputs(config, batch_size=batch_size)
|
||||
# # outputs = model.inference(input_dummy, aux_input={"x_lengths": input_lengths})
|
||||
# # self._check_inference_outputs(outputs, input_dummy, batch_size=batch_size)
|
||||
|
||||
# def test_multispeaker_inference(self):
|
||||
# num_speakers = 10
|
||||
# model_args = ForwardTTSE2EArgs(
|
||||
# spec_segment_size=10, num_speakers=num_speakers, use_speaker_embedding=True
|
||||
# )
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
|
||||
# batch_size = 1
|
||||
# input_dummy, *_ = self._create_inputs(config, batch_size=batch_size)
|
||||
# speaker_ids = torch.randint(0, num_speakers, (batch_size,)).long().to(device)
|
||||
# outputs = model.inference(input_dummy, {"speaker_ids": speaker_ids})
|
||||
# self._check_inference_outputs(outputs, input_dummy, batch_size=batch_size)
|
||||
|
||||
# # batch_size = 2
|
||||
# # input_dummy, input_lengths, *_ = self._create_inputs(config, batch_size=batch_size)
|
||||
# # speaker_ids = torch.randint(0, num_speakers, (batch_size,)).long().to(device)
|
||||
# # outputs = model.inference(input_dummy, {"x_lengths": input_lengths, "speaker_ids": speaker_ids})
|
||||
# # self._check_inference_outputs(outputs, input_dummy, batch_size=batch_size)
|
||||
|
||||
# # def test_multilingual_inference(self):
|
||||
# # """TODO"""
|
||||
|
||||
# def test_d_vector_inference(self):
|
||||
# model_args = ForwardTTSE2EArgs(
|
||||
# spec_segment_size=10,
|
||||
# num_chars=32,
|
||||
# use_d_vector_file=True,
|
||||
# d_vector_dim=256,
|
||||
# d_vector_file=os.path.join(get_tests_data_path(), "dummy_speakers.json"),
|
||||
# )
|
||||
# config = FastPitchE2EConfig(model_args=model_args)
|
||||
# model = ForwardTTSE2E(config).to(device)
|
||||
# model.eval()
|
||||
# # batch size = 1
|
||||
# input_dummy = torch.randint(0, 24, (1, 128)).long().to(device)
|
||||
# d_vectors = torch.randn(1, 256).to(device)
|
||||
# outputs = model.inference(input_dummy, aux_input={"d_vectors": d_vectors})
|
||||
# self._check_inference_outputs(outputs, input_dummy)
|
||||
# # batch size = 2
|
||||
# # input_dummy, input_lengths, *_ = self._create_inputs(config)
|
||||
# # d_vectors = torch.randn(2, 256).to(device)
|
||||
# # outputs = model.inference(input_dummy, aux_input={"x_lengths": input_lengths, "d_vectors": d_vectors})
|
||||
# # self._check_inference_outputs(outputs, input_dummy, batch_size=2)
|
||||
|
||||
def test_train_step(self):
|
||||
# setup the model
|
||||
with torch.autograd.set_detect_anomaly(True):
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E(config).to(device)
|
||||
model.train()
|
||||
# model to train
|
||||
optimizers = model.get_optimizer()
|
||||
criterions = model.get_criterion()
|
||||
criterions = [criterions[0].to(device), criterions[1].to(device)]
|
||||
# reference model to compare model weights
|
||||
model_ref = ForwardTTSE2E(config).to(device)
|
||||
# # pass the state to ref model
|
||||
model_ref.load_state_dict(copy.deepcopy(model.state_dict()))
|
||||
count = 0
|
||||
for param, param_ref in zip(model.parameters(), model_ref.parameters()):
|
||||
assert (param - param_ref).sum() == 0, param
|
||||
count = count + 1
|
||||
for _ in range(5):
|
||||
batch = self._create_batch(config, 2)
|
||||
for idx in [0, 1]:
|
||||
outputs, loss_dict = model.train_step(batch, criterions, idx)
|
||||
self.assertFalse(not outputs)
|
||||
self.assertFalse(not loss_dict)
|
||||
loss_dict["loss"].backward()
|
||||
optimizers[idx].step()
|
||||
optimizers[idx].zero_grad()
|
||||
|
||||
# check parameter changes
|
||||
self._check_parameter_changes(model, model_ref)
|
||||
|
||||
def test_train_eval_log(self):
|
||||
batch_size = 2
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
model.train()
|
||||
batch = self._create_batch(config, batch_size)
|
||||
logger = TensorboardLogger(
|
||||
log_dir=os.path.join(get_tests_output_path(), "dummy_fast_pitch_e2e_logs"),
|
||||
model_name="fast_pitch_e2e_test_train_log",
|
||||
)
|
||||
criterion = model.get_criterion()
|
||||
criterion = [criterion[0].to(device), criterion[1].to(device)]
|
||||
outputs = [None] * 2
|
||||
outputs[0], _ = model.train_step(batch, criterion, 0)
|
||||
outputs[1], _ = model.train_step(batch, criterion, 1)
|
||||
model.train_log(batch=batch, outputs=outputs, logger=logger, assets=None, steps=1)
|
||||
model.eval_log(batch, outputs, logger, None, 1)
|
||||
logger.finish()
|
||||
|
||||
def test_test_run(self):
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
model.eval()
|
||||
test_figures, test_audios = model.test_run(None)
|
||||
self.assertTrue(test_figures is not None)
|
||||
self.assertTrue(test_audios is not None)
|
||||
|
||||
def test_load_checkpoint(self):
|
||||
chkp_path = os.path.join(get_tests_output_path(), "dummy_fast_pitch_e2e_tts_checkpoint.pth")
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
chkp = {}
|
||||
chkp["model"] = model.state_dict()
|
||||
torch.save(chkp, chkp_path)
|
||||
model.load_checkpoint(config, chkp_path)
|
||||
self.assertTrue(model.training)
|
||||
model.load_checkpoint(config, chkp_path, eval=True)
|
||||
self.assertFalse(model.training)
|
||||
|
||||
def test_get_criterion(self):
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
criterion = model.get_criterion()
|
||||
self.assertTrue(criterion is not None)
|
||||
|
||||
def test_init_from_config(self):
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
|
||||
model_args = ForwardTTSE2EArgs(spec_segment_size=10, num_speakers=2)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
self.assertTrue(not hasattr(model, "emb_g"))
|
||||
|
||||
model_args = ForwardTTSE2EArgs(
|
||||
spec_segment_size=10, num_speakers=2, use_speaker_embedding=True
|
||||
)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
self.assertEqual(model.num_speakers, 2)
|
||||
self.assertTrue(hasattr(model, "emb_g"))
|
||||
|
||||
model_args = ForwardTTSE2EArgs(
|
||||
spec_segment_size=10,
|
||||
num_speakers=2,
|
||||
use_speaker_embedding=True,
|
||||
speakers_file=os.path.join(get_tests_data_path(), "ljspeech", "speakers.json"),
|
||||
)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
self.assertEqual(model.num_speakers, 10)
|
||||
self.assertTrue(hasattr(model, "emb_g"))
|
||||
|
||||
model_args = ForwardTTSE2EArgs(
|
||||
spec_segment_size=10,
|
||||
use_d_vector_file=True,
|
||||
d_vector_dim=256,
|
||||
d_vector_file=os.path.join(get_tests_data_path(), "ljspeech", "speakers.json"),
|
||||
)
|
||||
config = FastPitchE2EConfig(model_args=model_args)
|
||||
model = ForwardTTSE2E.init_from_config(config, verbose=False).to(device)
|
||||
self.assertTrue(model.num_speakers == 10)
|
||||
self.assertTrue(not hasattr(model, "emb_g"))
|
||||
self.assertTrue(model.embedded_speaker_dim == config.model_args.d_vector_dim)
|
||||
Loading…
Reference in New Issue