From 77d85c6cc5a3fb0cd5e852d0430468112c28a66c Mon Sep 17 00:00:00 2001 From: Edresson Date: Mon, 10 May 2021 17:08:38 -0300 Subject: [PATCH] add softmaxproto loss and bug fix in data loader --- TTS/bin/train_encoder.py | 47 ++++++----- TTS/speaker_encoder/{ => configs}/config.json | 3 + .../configs/config_softmaxproto.json | 78 +++++++++++++++++++ TTS/speaker_encoder/dataset.py | 58 +++++++++++--- TTS/speaker_encoder/losses.py | 64 ++++++++++++++- TTS/speaker_encoder/utils/generic_utils.py | 2 +- 6 files changed, 219 insertions(+), 33 deletions(-) rename TTS/speaker_encoder/{ => configs}/config.json (97%) create mode 100644 TTS/speaker_encoder/configs/config_softmaxproto.json diff --git a/TTS/bin/train_encoder.py b/TTS/bin/train_encoder.py index 3a3f876e..05a76b68 100644 --- a/TTS/bin/train_encoder.py +++ b/TTS/bin/train_encoder.py @@ -11,7 +11,7 @@ import torch from torch.utils.data import DataLoader from TTS.speaker_encoder.dataset import MyDataset -from TTS.speaker_encoder.losses import AngleProtoLoss, GE2ELoss +from TTS.speaker_encoder.losses import AngleProtoLoss, GE2ELoss, SoftmaxLoss, SoftmaxAngleProtoLoss from TTS.speaker_encoder.model import SpeakerEncoder from TTS.speaker_encoder.utils.generic_utils import check_config_speaker_encoder, save_best_model from TTS.speaker_encoder.utils.visual import plot_embeddings @@ -45,15 +45,16 @@ def setup_loader(ap: AudioProcessor, is_val: bool = False, verbose: bool = False dataset = MyDataset( ap, meta_data_eval if is_val else meta_data_train, - voice_len=1.6, + voice_len=getattr(c, "voice_len", 1.6), num_utter_per_speaker=c.num_utters_per_speaker, num_speakers_in_batch=c.num_speakers_in_batch, - skip_speakers=False, + skip_speakers=getattr(c, "skip_speakers", False), storage_size=c.storage["storage_size"], sample_from_storage_p=c.storage["sample_from_storage_p"], additive_noise=c.storage["additive_noise"], verbose=verbose, ) + # sampler = DistributedSampler(dataset) if num_gpus > 1 else None loader = DataLoader( dataset, @@ -62,11 +63,25 @@ def setup_loader(ap: AudioProcessor, is_val: bool = False, verbose: bool = False num_workers=c.num_loader_workers, collate_fn=dataset.collate_fn, ) - return loader + return loader, dataset.get_num_speakers() -def train(model, criterion, optimizer, scheduler, ap, global_step): - data_loader = setup_loader(ap, is_val=False, verbose=True) +def train(model, optimizer, scheduler, ap, global_step): + data_loader, num_speakers = setup_loader(ap, is_val=False, verbose=True) + + if c.loss == "ge2e": + criterion = GE2ELoss(loss_method="softmax") + elif c.loss == "angleproto": + criterion = AngleProtoLoss() + elif c.loss == "softmaxproto": + criterion = SoftmaxAngleProtoLoss(c.model["proj_dim"], num_speakers) + else: + raise Exception("The %s not is a loss supported" % c.loss) + + if use_cuda: + model = model.cuda() + criterion.cuda() + model.train() epoch_time = 0 best_loss = float("inf") @@ -77,7 +92,8 @@ def train(model, criterion, optimizer, scheduler, ap, global_step): start_time = time.time() # setup input data - inputs = data[0] + inputs, labels = data + loader_time = time.time() - end_time global_step += 1 @@ -89,13 +105,13 @@ def train(model, criterion, optimizer, scheduler, ap, global_step): # dispatch data to GPU if use_cuda: inputs = inputs.cuda(non_blocking=True) - # labels = labels.cuda(non_blocking=True) + labels = labels.cuda(non_blocking=True) # forward pass model outputs = model(inputs) # loss computation - loss = criterion(outputs.view(c.num_speakers_in_batch, outputs.shape[0] // c.num_speakers_in_batch, -1)) + loss = criterion(outputs.view(c.num_speakers_in_batch, outputs.shape[0] // c.num_speakers_in_batch, -1), labels) loss.backward() grad_norm, _ = check_update(model, c.grad_clip) optimizer.step() @@ -158,13 +174,6 @@ def main(args): # pylint: disable=redefined-outer-name ) optimizer = RAdam(model.parameters(), lr=c.lr) - if c.loss == "ge2e": - criterion = GE2ELoss(loss_method="softmax") - elif c.loss == "angleproto": - criterion = AngleProtoLoss() - else: - raise Exception("The %s not is a loss supported" % c.loss) - if args.restore_path: checkpoint = torch.load(args.restore_path) try: @@ -187,10 +196,6 @@ def main(args): # pylint: disable=redefined-outer-name else: args.restore_step = 0 - if use_cuda: - model = model.cuda() - criterion.cuda() - if c.lr_decay: scheduler = NoamLR(optimizer, warmup_steps=c.warmup_steps, last_epoch=args.restore_step - 1) else: @@ -203,7 +208,7 @@ def main(args): # pylint: disable=redefined-outer-name meta_data_train, meta_data_eval = load_meta_data(c.datasets) global_step = args.restore_step - _, global_step = train(model, criterion, optimizer, scheduler, ap, global_step) + _, global_step = train(model, optimizer, scheduler, ap, global_step) if __name__ == "__main__": diff --git a/TTS/speaker_encoder/config.json b/TTS/speaker_encoder/configs/config.json similarity index 97% rename from TTS/speaker_encoder/config.json rename to TTS/speaker_encoder/configs/config.json index 4fbd84cc..2b437e5a 100644 --- a/TTS/speaker_encoder/config.json +++ b/TTS/speaker_encoder/configs/config.json @@ -37,6 +37,9 @@ "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, // + "skip_speakers": false, // skip speakers with samples less than "num_utters_per_speaker" + + "voice_len": 1.6, // number of seconds for each training instance "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" diff --git a/TTS/speaker_encoder/configs/config_softmaxproto.json b/TTS/speaker_encoder/configs/config_softmaxproto.json new file mode 100644 index 00000000..f7a24a15 --- /dev/null +++ b/TTS/speaker_encoder/configs/config_softmaxproto.json @@ -0,0 +1,78 @@ + +{ + "run_name": "speaker_encoder", + "run_description": "train speaker encoder with VCTK", + "audio":{ + // Audio processing parameters + "num_mels": 80, // size of the mel spec frame. + "fft_size": 1024, // number of stft frequency levels. Size of the linear spectogram frame. + "sample_rate": 16000, // DATASET-RELATED: wav sample-rate. If different than the original data, it is resampled. + "win_length": 1024, // stft window length in ms. + "hop_length": 256, // stft window hop-lengh in ms. + "frame_length_ms": null, // stft window length in ms.If null, 'win_length' is used. + "frame_shift_ms": null, // stft window hop-lengh in ms. If null, 'hop_length' is used. + "preemphasis": 0.98, // pre-emphasis to reduce spec noise and make it more structured. If 0.0, no -pre-emphasis. + "min_level_db": -100, // normalization range + "ref_level_db": 20, // reference level db, theoretically 20db is the sound of air. + "power": 1.5, // value to sharpen wav signals after GL algorithm. + "griffin_lim_iters": 60,// #griffin-lim iterations. 30-60 is a good range. Larger the value, slower the generation. + "stft_pad_mode": "reflect", + // Normalization parameters + "signal_norm": true, // normalize the spec values in range [0, 1] + "symmetric_norm": true, // move normalization to range [-1, 1] + "max_norm": 4.0, // scale normalization to range [-max_norm, max_norm] or [0, max_norm] + "clip_norm": true, // clip normalized values into the range. + "mel_fmin": 0.0, // minimum freq level for mel-spec. ~50 for male and ~95 for female voices. Tune for dataset!! + "mel_fmax": 8000.0, // maximum freq level for mel-spec. Tune for dataset!! + "spec_gain": 20.0, + "do_trim_silence": false, // enable trimming of slience of audio as you load it. LJspeech (false), TWEB (false), Nancy (true) + "trim_db": 60, // threshold for timming silence. Set this according to your dataset. + "stats_path": null // DO NOT USE WITH MULTI_SPEAKER MODEL. scaler stats file computed by 'compute_statistics.py'. If it is defined, mean-std based notmalization is used and other normalization params are ignored + }, + "reinit_layers": [], + + "loss": "softmaxproto", // "ge2e" to use Generalized End-to-End loss, "angleproto" to use Angular Prototypical loss and "softmaxproto" to use Softmax with Angular Prototypical loss + "grad_clip": 3.0, // upper limit for gradients for clipping. + "epochs": 1000, // total number of epochs to train. + "lr": 0.0001, // Initial learning rate. If Noam decay is active, maximum learning rate. + "lr_decay": false, // if true, Noam learning rate decaying is applied through training. + "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. + + // Speakers config + "num_speakers_in_batch": 108, // Batch size for training. + "num_utters_per_speaker": 2, // + "skip_speakers": true, // skip speakers with samples less than "num_utters_per_speaker" + + "voice_len": 2, // number of seconds for each training instance + + "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" + "save_step": 1000, // Number of training steps expected to save traning stats and checkpoints. + "print_step": 20, // Number of steps to log traning on console. + "output_path": "../../../checkpoints/speaker_encoder/", // DATASET-RELATED: output path for all training outputs. + + "model": { + "input_dim": 80, + "proj_dim": 512, + "lstm_dim": 768, + "num_lstm_layers": 3, + "use_lstm_with_projection": true + }, + "storage": { + "sample_from_storage_p": 0.66, // the probability with which we'll sample from the DataSet in-memory storage + "storage_size": 15, // the size of the in-memory storage with respect to a single batch + "additive_noise": 1e-5 // add very small gaussian noise to the data in order to increase robustness + }, + "datasets": + [ + { + "name": "vctk", + "path": "/workspace/store/ecasanova/datasets/VCTK-Corpus-removed-silence/", + "meta_file_train": null, + "meta_file_val": null + } + ] +} \ No newline at end of file diff --git a/TTS/speaker_encoder/dataset.py b/TTS/speaker_encoder/dataset.py index 38d8b5f9..3308fef9 100644 --- a/TTS/speaker_encoder/dataset.py +++ b/TTS/speaker_encoder/dataset.py @@ -30,7 +30,6 @@ class MyDataset(Dataset): super().__init__() self.items = meta_data self.sample_rate = ap.sample_rate - self.voice_len = voice_len self.seq_len = int(voice_len * self.sample_rate) self.num_speakers_in_batch = num_speakers_in_batch self.num_utter_per_speaker = num_utter_per_speaker @@ -41,10 +40,15 @@ class MyDataset(Dataset): self.storage = queue.Queue(maxsize=storage_size * num_speakers_in_batch) self.sample_from_storage_p = float(sample_from_storage_p) self.additive_noise = float(additive_noise) + + speakers_aux = list(self.speakers) + speakers_aux.sort() + self.speakerid_to_classid = {key : i for i, key in enumerate(speakers_aux)} + if self.verbose: print("\n > DataLoader initialization") print(f" | > Speakers per Batch: {num_speakers_in_batch}") - print(f" | > Storage Size: {self.storage.maxsize} speakers, each with {num_utter_per_speaker} utters") + print(f" | > Storage Size: {self.storage.maxsize} instances, each with {num_utter_per_speaker} utters") print(f" | > Sample_from_storage_p : {self.sample_from_storage_p}") print(f" | > Noise added : {self.additive_noise}") print(f" | > Number of instances : {len(self.items)}") @@ -110,8 +114,16 @@ class MyDataset(Dataset): def __len__(self): return int(1e10) - def __sample_speaker(self): + def get_num_speakers(self): + return len(self.speakers) + + def __sample_speaker(self, ignore_speakers=None): speaker = random.sample(self.speakers, 1)[0] + # if list of speakers_id is provide make sure that it's will be ignored + if ignore_speakers: + while self.speakerid_to_classid[speaker] in ignore_speakers: + speaker = random.sample(self.speakers, 1)[0] + if self.num_utter_per_speaker > len(self.speaker_to_utters[speaker]): utters = random.choices(self.speaker_to_utters[speaker], k=self.num_utter_per_speaker) else: @@ -127,7 +139,8 @@ class MyDataset(Dataset): for _ in range(self.num_utter_per_speaker): # TODO:dummy but works while True: - if len(self.speaker_to_utters[speaker]) > 0: + # remove speakers that have num_utter less than 2 + if len(self.speaker_to_utters[speaker]) > 1: utter = random.sample(self.speaker_to_utters[speaker], 1)[0] else: self.speakers.remove(speaker) @@ -139,21 +152,47 @@ class MyDataset(Dataset): self.speaker_to_utters[speaker].remove(utter) wavs.append(wav) - labels.append(speaker) + labels.append(self.speakerid_to_classid[speaker]) return wavs, labels def __getitem__(self, idx): speaker, _ = self.__sample_speaker() - return speaker + speaker_id = self.speakerid_to_classid[speaker] + return speaker, speaker_id def collate_fn(self, batch): + # get the batch speaker_ids + batch = np.array(batch) + speakers_id_in_batch = set(batch[:, 1].astype(np.int32)) + labels = [] feats = [] - for speaker in batch: + speakers = set() + for speaker, speaker_id in batch: + if random.random() < self.sample_from_storage_p and self.storage.full(): # sample from storage (if full), ignoring the speaker wavs_, labels_ = random.choice(self.storage.queue) + + # force choose the current speaker or other not in batch + '''while labels_[0] in speakers_id_in_batch: + if labels_[0] == speaker_id: + break + wavs_, labels_ = random.choice(self.storage.queue)''' + + speakers.add(labels_[0]) + speakers_id_in_batch.add(labels_[0]) + else: + # ensure that an speaker appears only once in the batch + if speaker_id in speakers: + speaker, _ = self.__sample_speaker(speakers_id_in_batch) + speaker_id = self.speakerid_to_classid[speaker] + # append the new speaker from batch + speakers_id_in_batch.add(speaker_id) + + speakers.add(speaker_id) + # don't sample from storage, but from HDD wavs_, labels_ = self.__sample_speaker_utterances(speaker) # if storage is full, remove an item @@ -167,14 +206,15 @@ class MyDataset(Dataset): noises_ = [np.random.normal(0, self.additive_noise, size=len(w)) for w in wavs_] wavs_ = [wavs_[i] + noises_[i] for i in range(len(wavs_))] - # get a random subset of each of the wavs and convert to MFCC. + # get a random subset of each of the wavs and extract mel spectrograms. offsets_ = [random.randint(0, wav.shape[0] - self.seq_len) for wav in wavs_] mels_ = [ self.ap.melspectrogram(wavs_[i][offsets_[i] : offsets_[i] + self.seq_len]) for i in range(len(wavs_)) ] feats_ = [torch.FloatTensor(mel) for mel in mels_] - labels.append(labels_) + labels.append(torch.LongTensor(labels_)) feats.extend(feats_) feats = torch.stack(feats) + labels = torch.stack(labels) return feats.transpose(1, 2), labels diff --git a/TTS/speaker_encoder/losses.py b/TTS/speaker_encoder/losses.py index 69264ab4..52871fb4 100644 --- a/TTS/speaker_encoder/losses.py +++ b/TTS/speaker_encoder/losses.py @@ -103,10 +103,13 @@ class GE2ELoss(nn.Module): L.append(L_row) return torch.stack(L) - def forward(self, dvecs): + def forward(self, dvecs, label=None): """ Calculates the GE2E loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) """ + + assert x.size()[1] >= 2 + centroids = torch.mean(dvecs, 1) cos_sim_matrix = self.calc_cosine_sim(dvecs, centroids) torch.clamp(self.w, 1e-6) @@ -138,10 +141,13 @@ class AngleProtoLoss(nn.Module): print(" > Initialised Angular Prototypical loss") - def forward(self, x): + def forward(self, x, label=None): """ Calculates the AngleProto loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) """ + + assert x.size()[1] >= 2 + out_anchor = torch.mean(x[:, 1:, :], 1) out_positive = x[:, 0, :] num_speakers = out_anchor.size()[0] @@ -155,3 +161,57 @@ class AngleProtoLoss(nn.Module): label = torch.arange(num_speakers).to(cos_sim_matrix.device) L = self.criterion(cos_sim_matrix, label) return L + +class SoftmaxLoss(nn.Module): + """ + Implementation of the Softmax loss as defined in https://arxiv.org/abs/2003.11982 + Args: + - embedding_dim (float): speaker embedding dim + - n_speakers (float): number of speakers + """ + def __init__(self, embedding_dim, n_speakers): + super().__init__() + + self.criterion = torch.nn.CrossEntropyLoss() + self.fc = nn.Linear(embedding_dim, n_speakers) + + print('Initialised Softmax Loss') + + def forward(self, x, label=None): + + x = self.fc(x) + L = self.criterion(x, label) + + return L + +class SoftmaxAngleProtoLoss(nn.Module): + """ + Implementation of the Softmax AnglePrototypical loss as defined in https://arxiv.org/abs/2009.14153 + Args: + - embedding_dim (float): speaker embedding dim + - n_speakers (float): number of speakers + - init_w (float): defines the initial value of w + - init_b (float): definies the initial value of b + """ + def __init__(self, embedding_dim, n_speakers, init_w=10.0, init_b=-5.0): + super().__init__() + + self.softmax = SoftmaxLoss(embedding_dim, n_speakers) + self.angleproto = AngleProtoLoss(init_w, init_b) + + print('Initialised SoftmaxAnglePrototypical Loss') + + def forward(self, x, label=None): + """ + Calculates the SoftmaxAnglePrototypical loss for an input of dimensions (num_speakers, num_utts_per_speaker, dvec_feats) + """ + + assert x.size()[1] == 2 + + Lp = self.angleproto(x) + + x = x.reshape(-1, x.size()[-1]) + label = label.reshape(-1) + Ls = self.softmax(x, label) + + return Ls+Lp diff --git a/TTS/speaker_encoder/utils/generic_utils.py b/TTS/speaker_encoder/utils/generic_utils.py index c9bfa679..c50304cc 100644 --- a/TTS/speaker_encoder/utils/generic_utils.py +++ b/TTS/speaker_encoder/utils/generic_utils.py @@ -82,7 +82,7 @@ def check_config_speaker_encoder(c): check_argument("griffin_lim_iters", c["audio"], restricted=True, val_type=int, min_val=10, max_val=1000) # training parameters - check_argument("loss", c, enum_list=["ge2e", "angleproto"], restricted=True, val_type=str) + check_argument("loss", c, enum_list=["ge2e", "angleproto", "softmaxproto"], restricted=True, val_type=str) check_argument("grad_clip", c, restricted=True, val_type=float) check_argument("epochs", c, restricted=True, val_type=int, min_val=1) check_argument("lr", c, restricted=True, val_type=float, min_val=0)