REBASED: Transform Speaker Encoder in a Generic Encoder and Implement Emotion Encoder training support (#1349)

* Rename Speaker encoder module to encoder

* Add a generic emotion dataset formatter

* Transform the Speaker Encoder dataset to a generic dataset and create emotion encoder config

* Add class map in emotion config

* Add Base encoder config

* Add evaluation encoder script

* Fix the bug in plot_embeddings

* Enable Weight decay for encoder training

* Add argumnet to disable storage

* Add Perfect Sampler and remove storage

* Add evaluation during encoder training

* Fix lint checks

* Remove useless config parameter

* Active evaluation in speaker encoder test and use multispeaker dataset for this test

* Unit tests fixs

* Remove useless tests for speedup the aux_tests

* Use get_optimizer in Encoder

* Add BaseEncoder Class

* Fix the unitests

* Add Perfect Batch Sampler unit test

* Add compute encoder accuracy in a function

tests/aux_tests/test_speaker_encoder.py

@@ -3,9 +3,9 @@ import unittest
 import torch as T
 
 from tests import get_tests_input_path
-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.encoder.losses import AngleProtoLoss, GE2ELoss, SoftmaxAngleProtoLoss
+from TTS.encoder.models.lstm import LSTMSpeakerEncoder
+from TTS.encoder.models.resnet import ResNetSpeakerEncoder
 
 file_path = get_tests_input_path()
 

tests/aux_tests/test_speaker_encoder_train.py

@@ -4,14 +4,14 @@ import shutil
 
 from tests import get_device_id, get_tests_output_path, run_cli
 from TTS.config.shared_configs import BaseAudioConfig
-from TTS.speaker_encoder.speaker_encoder_config import SpeakerEncoderConfig
+from TTS.encoder.configs.speaker_encoder_config import SpeakerEncoderConfig
 
 
 def run_test_train():
     command = (
         f"CUDA_VISIBLE_DEVICES='{get_device_id()}' python TTS/bin/train_encoder.py --config_path {config_path} "
         f"--coqpit.output_path {output_path} "
-        "--coqpit.datasets.0.name ljspeech "
+        "--coqpit.datasets.0.name ljspeech_test "
         "--coqpit.datasets.0.meta_file_train metadata.csv "
         "--coqpit.datasets.0.meta_file_val metadata.csv "
         "--coqpit.datasets.0.path tests/data/ljspeech "
@@ -24,17 +24,21 @@ output_path = os.path.join(get_tests_output_path(), "train_outputs")
 
 config = SpeakerEncoderConfig(
     batch_size=4,
-    num_speakers_in_batch=1,
-    num_utters_per_speaker=10,
-    num_loader_workers=0,
-    max_train_step=2,
+    num_classes_in_batch=4,
+    num_utter_per_class=2,
+    eval_num_classes_in_batch=4,
+    eval_num_utter_per_class=2,
+    num_loader_workers=1,
+    epochs=1,
     print_step=1,
-    save_step=1,
+    save_step=2,
     print_eval=True,
+    run_eval=True,
     audio=BaseAudioConfig(num_mels=80),
 )
 config.audio.do_trim_silence = True
 config.audio.trim_db = 60
+config.loss = "ge2e"
 config.save_json(config_path)
 
 print(config)
@@ -69,14 +73,14 @@ run_cli(command_train)
 shutil.rmtree(continue_path)
 
 # test model with ge2e loss function
-config.loss = "ge2e"
-config.save_json(config_path)
-run_test_train()
+# config.loss = "ge2e"
+# config.save_json(config_path)
+# run_test_train()
 
 # test model with angleproto loss function
-config.loss = "angleproto"
-config.save_json(config_path)
-run_test_train()
+# config.loss = "angleproto"
+# config.save_json(config_path)
+# run_test_train()
 
 # test model with softmaxproto loss function
 config.loss = "softmaxproto"

tests/aux_tests/test_speaker_manager.py

@@ -6,8 +6,8 @@ import torch
 
 from tests import get_tests_input_path
 from TTS.config import load_config
-from TTS.speaker_encoder.utils.generic_utils import setup_speaker_encoder_model
-from TTS.speaker_encoder.utils.io import save_checkpoint
+from TTS.encoder.utils.generic_utils import setup_speaker_encoder_model
+from TTS.encoder.utils.io import save_checkpoint
 from TTS.tts.utils.speakers import SpeakerManager
 from TTS.utils.audio import AudioProcessor
 

tests/data_tests/test_samplers.py

@@ -8,6 +8,7 @@ from TTS.config.shared_configs import BaseDatasetConfig
 from TTS.tts.datasets import load_tts_samples
 from TTS.tts.utils.languages import get_language_balancer_weights
 from TTS.tts.utils.speakers import get_speaker_balancer_weights
+from TTS.encoder.utils.samplers import PerfectBatchSampler
 
 # Fixing random state to avoid random fails
 torch.manual_seed(0)
@@ -82,3 +83,51 @@ class TestSamplers(unittest.TestCase):
                 spk2 += 1
 
         assert is_balanced(spk1, spk2), "Speaker Weighted sampler is supposed to be balanced"
+
+    def test_perfect_sampler(self):  # pylint: disable=no-self-use
+        classes = set()
+        for item in train_samples:
+            classes.add(item["speaker_name"])
+
+        sampler = PerfectBatchSampler(
+            train_samples,
+            classes,
+            batch_size=2 * 3, # total batch size
+            num_classes_in_batch=2,
+            label_key="speaker_name",
+            shuffle=False,
+            drop_last=True)
+        batchs = functools.reduce(lambda a, b: a + b, [list(sampler) for i in range(100)])
+        for batch in batchs:
+            spk1, spk2 = 0, 0
+            # for in each batch
+            for index in batch:
+                if train_samples[index]["speaker_name"] == "ljspeech-0":
+                    spk1 += 1
+                else:
+                    spk2 += 1
+            assert spk1 == spk2, "PerfectBatchSampler is supposed to be perfectly balanced"
+
+    def test_perfect_sampler_shuffle(self):  # pylint: disable=no-self-use
+        classes = set()
+        for item in train_samples:
+            classes.add(item["speaker_name"])
+
+        sampler = PerfectBatchSampler(
+            train_samples,
+            classes,
+            batch_size=2 * 3, # total batch size
+            num_classes_in_batch=2,
+            label_key="speaker_name",
+            shuffle=True,
+            drop_last=False)
+        batchs = functools.reduce(lambda a, b: a + b, [list(sampler) for i in range(100)])
+        for batch in batchs:
+            spk1, spk2 = 0, 0
+            # for in each batch
+            for index in batch:
+                if train_samples[index]["speaker_name"] == "ljspeech-0":
+                    spk1 += 1
+                else:
+                    spk2 += 1
+            assert spk1 == spk2, "PerfectBatchSampler is supposed to be perfectly balanced"

tests/inputs/test_glow_tts.json

@@ -66,8 +66,8 @@
     "use_mas": false,       // use Monotonic Alignment Search if true. Otherwise use pre-computed attention alignments.
 
     // TRAINING
-    "batch_size": 2,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
-    "eval_batch_size":1,
+    "batch_size": 8,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
+    "eval_batch_size": 8,
     "r": 1,                 // Number of decoder frames to predict per iteration. Set the initial values if gradual training is enabled.
     "loss_masking": true,   // enable / disable loss masking against the sequence padding.
     "data_dep_init_iter": 1,

tests/inputs/test_speaker_encoder_config.json

@@ -36,8 +36,8 @@
     "warmup_steps": 4000, // Noam decay steps to increase the learning rate from 0 to "lr"
     "tb_model_param_stats": false, // true, plots param stats per layer on tensorboard. Might be memory consuming, but good for debugging.
     "steps_plot_stats": 10, // number of steps to plot embeddings.
-    "num_speakers_in_batch": 64, // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
-    "num_utters_per_speaker": 10,  //
+    "num_classes_in_batch": 64, // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
+    "num_utter_per_class": 10,  //
     "num_loader_workers": 8,        // number of training data loader processes. Don't set it too big. 4-8 are good values.
     "wd": 0.000001, // Weight decay weight.
     "checkpoint": true, // If true, it saves checkpoints per "save_step"

tests/inputs/test_tacotron2_config.json

@@ -61,8 +61,8 @@
     "reinit_layers": [],    // give a list of layer names to restore from the given checkpoint. If not defined, it reloads all heuristically matching layers.
 
     // TRAINING
-    "batch_size": 1,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
-    "eval_batch_size":1,
+    "batch_size": 8,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
+    "eval_batch_size": 8,
     "r": 7,                 // Number of decoder frames to predict per iteration. Set the initial values if gradual training is enabled.
     "gradual_training": [[0, 7, 4], [1, 5, 2]], //set gradual training steps [first_step, r, batch_size]. If it is null, gradual training is disabled. For Tacotron, you might need to reduce the 'batch_size' as you proceeed.
     "loss_masking": true,         // enable / disable loss masking against the sequence padding.

tests/inputs/test_tacotron_config.json

@@ -61,8 +61,8 @@
     "reinit_layers": [],    // give a list of layer names to restore from the given checkpoint. If not defined, it reloads all heuristically matching layers.
 
     // TRAINING
-    "batch_size": 1,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
-    "eval_batch_size":1,
+    "batch_size": 8,       // Batch size for training. Lower values than 32 might cause hard to learn attention. It is overwritten by 'gradual_training'.
+    "eval_batch_size": 8,
     "r": 7,                 // Number of decoder frames to predict per iteration. Set the initial values if gradual training is enabled.
     "gradual_training": [[0, 7, 4], [1, 5, 2]], //set gradual training steps [first_step, r, batch_size]. If it is null, gradual training is disabled. For Tacotron, you might need to reduce the 'batch_size' as you proceeed.
     "loss_masking": true,         // enable / disable loss masking against the sequence padding.

tests/tts_tests/test_vits.py

@@ -7,7 +7,7 @@ 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.config import load_config
-from TTS.speaker_encoder.utils.generic_utils import setup_speaker_encoder_model
+from TTS.encoder.utils.generic_utils import setup_speaker_encoder_model
 from TTS.tts.configs.vits_config import VitsConfig
 from TTS.tts.models.vits import Vits, VitsArgs, amp_to_db, db_to_amp, load_audio, spec_to_mel, wav_to_mel, wav_to_spec
 from TTS.tts.utils.speakers import SpeakerManager