From 0e372e0b9b4d067ee29f02915b4134c17f19f648 Mon Sep 17 00:00:00 2001 From: Edresson Casanova Date: Fri, 4 Mar 2022 15:09:51 -0300 Subject: [PATCH] Add Perfect Sampler and remove storage --- TTS/bin/eval_encoder.py | 2 +- TTS/bin/train_encoder.py | 46 +++-- TTS/encoder/configs/base_encoder_config.py | 8 - TTS/encoder/dataset.py | 217 +++++---------------- TTS/encoder/utils/samplers.py | 100 ++++++++++ 5 files changed, 189 insertions(+), 184 deletions(-) create mode 100644 TTS/encoder/utils/samplers.py diff --git a/TTS/bin/eval_encoder.py b/TTS/bin/eval_encoder.py index 0b1af9f2..9a4e0204 100644 --- a/TTS/bin/eval_encoder.py +++ b/TTS/bin/eval_encoder.py @@ -86,4 +86,4 @@ for key in class_acc_dict: print("Class", key, "Accuracy:", acc) acc_avg += acc -print("Average Accuracy:", acc_avg/len(class_acc_dict)) \ No newline at end of file +print("Average Accuracy:", acc_avg/len(class_acc_dict)) diff --git a/TTS/bin/train_encoder.py b/TTS/bin/train_encoder.py index c65474db..62cfbc71 100644 --- a/TTS/bin/train_encoder.py +++ b/TTS/bin/train_encoder.py @@ -13,6 +13,7 @@ from trainer.torch import NoamLR from TTS.encoder.dataset import EncoderDataset from TTS.encoder.losses import AngleProtoLoss, GE2ELoss, SoftmaxAngleProtoLoss from TTS.encoder.utils.generic_utils import save_best_model, setup_speaker_encoder_model +from TTS.encoder.utils.samplers import PerfectBatchSampler from TTS.encoder.utils.training import init_training from TTS.encoder.utils.visual import plot_embeddings from TTS.tts.datasets import load_tts_samples @@ -41,21 +42,24 @@ def setup_loader(ap: AudioProcessor, is_val: bool = False, verbose: bool = False voice_len=c.voice_len, num_utter_per_class=c.num_utter_per_class, num_classes_in_batch=c.num_classes_in_batch, - use_storage=c.use_storage, - skip_classes=c.skip_classes, - storage_size=c.storage["storage_size"], - sample_from_storage_p=c.storage["sample_from_storage_p"], verbose=verbose, augmentation_config=c.audio_augmentation if not is_val else None, use_torch_spec=c.model_params.get("use_torch_spec", False), ) - # sampler = DistributedSampler(dataset) if num_gpus > 1 else None + sampler = PerfectBatchSampler( + dataset.items, + dataset.get_class_list(), + batch_size=c.num_classes_in_batch*c.num_utter_per_class, # total batch size + num_classes_in_batch=c.num_classes_in_batch, + num_gpus=1, + shuffle=False if is_val else True, + drop_last=True) + loader = DataLoader( dataset, - batch_size=c.num_classes_in_batch, - shuffle=False, num_workers=c.num_loader_workers, + batch_sampler=sampler, collate_fn=dataset.collate_fn, ) @@ -70,12 +74,31 @@ def train(model, optimizer, scheduler, criterion, data_loader, global_step): avg_loss_all = 0 avg_loader_time = 0 end_time = time.time() - + print(len(data_loader)) for _, data in enumerate(data_loader): start_time = time.time() # setup input data inputs, labels = data + # agroup samples of each class in the batch. perfect sampler produces [3,2,1,3,2,1] we need [3,3,2,2,1,1] + labels = torch.transpose(labels.view(c.num_utter_per_class, c.num_classes_in_batch), 0, 1).reshape(labels.shape) + inputs = torch.transpose(inputs.view(c.num_utter_per_class, c.num_classes_in_batch, -1), 0, 1).reshape(inputs.shape) + """ + labels_converted = torch.transpose(labels.view(c.num_utter_per_class, c.num_classes_in_batch), 0, 1).reshape(labels.shape) + inputs_converted = torch.transpose(inputs.view(c.num_utter_per_class, c.num_classes_in_batch, -1), 0, 1).reshape(inputs.shape) + idx = 0 + for j in range(0, c.num_classes_in_batch, 1): + for i in range(j, len(labels), c.num_classes_in_batch): + if not torch.all(labels[i].eq(labels_converted[idx])) or not torch.all(inputs[i].eq(inputs_converted[idx])): + print("Invalid") + print(labels) + exit() + idx += 1 + labels = labels_converted + inputs = inputs_converted + print(labels) + print(inputs.shape)""" + loader_time = time.time() - end_time global_step += 1 @@ -159,9 +182,10 @@ def main(args): # pylint: disable=redefined-outer-name optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=c.wd) # pylint: disable=redefined-outer-name - meta_data_train, meta_data_eval = load_tts_samples(c.datasets, eval_split=False) + meta_data_train, meta_data_eval = load_tts_samples(c.datasets, eval_split=True) - data_loader, num_classes, map_classid_to_classname = setup_loader(ap, is_val=False, verbose=True) + train_data_loader, num_classes, map_classid_to_classname = setup_loader(ap, is_val=False, verbose=True) + # eval_data_loader, _, _ = setup_loader(ap, is_val=True, verbose=True) if c.loss == "ge2e": criterion = GE2ELoss(loss_method="softmax") @@ -211,7 +235,7 @@ def main(args): # pylint: disable=redefined-outer-name criterion.cuda() global_step = args.restore_step - _, global_step = train(model, optimizer, scheduler, criterion, data_loader, global_step) + _, global_step = train(model, optimizer, scheduler, criterion, train_data_loader, global_step) if __name__ == "__main__": diff --git a/TTS/encoder/configs/base_encoder_config.py b/TTS/encoder/configs/base_encoder_config.py index 838f9300..b72e6076 100644 --- a/TTS/encoder/configs/base_encoder_config.py +++ b/TTS/encoder/configs/base_encoder_config.py @@ -27,14 +27,6 @@ class BaseEncoderConfig(BaseTrainingConfig): audio_augmentation: Dict = field(default_factory=lambda: {}) - use_storage: bool = False - storage: Dict = field( - default_factory=lambda: { - "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 - } - ) - # training params max_train_step: int = 1000000 # end training when number of training steps reaches this value. loss: str = "angleproto" diff --git a/TTS/encoder/dataset.py b/TTS/encoder/dataset.py index 474aa0c2..6b71c103 100644 --- a/TTS/encoder/dataset.py +++ b/TTS/encoder/dataset.py @@ -6,7 +6,6 @@ from torch.utils.data import Dataset from TTS.encoder.utils.generic_utils import AugmentWAV, Storage - class EncoderDataset(Dataset): def __init__( self, @@ -14,11 +13,7 @@ class EncoderDataset(Dataset): meta_data, voice_len=1.6, num_classes_in_batch=64, - use_storage=False, - storage_size=1, - sample_from_storage_p=0.5, num_utter_per_class=10, - skip_classes=False, verbose=False, augmentation_config=None, use_torch_spec=None, @@ -34,30 +29,15 @@ class EncoderDataset(Dataset): self.items = meta_data self.sample_rate = ap.sample_rate self.seq_len = int(voice_len * self.sample_rate) - self.num_classes_in_batch = num_classes_in_batch self.num_utter_per_class = num_utter_per_class - self.skip_classes = skip_classes - self.use_storage = use_storage self.ap = ap self.verbose = verbose self.use_torch_spec = use_torch_spec self.__parse_items() - storage_max_size = storage_size * num_classes_in_batch - if self.use_storage: - self.storage = Storage( - maxsize=storage_max_size, storage_batchs=storage_size, num_classes_in_batch=num_classes_in_batch - ) - self.sample_from_storage_p = float(sample_from_storage_p) - else: - self.storage = None - self.sample_from_storage_p = None + self.classname_to_classid = {key: i for i, key in enumerate(self.classes)} - classes_aux = list(self.classes) - classes_aux.sort() - self.classname_to_classid = {key: i for i, key in enumerate(classes_aux)} - - # Augmentation + # Data Augmentation self.augmentator = None self.gaussian_augmentation_config = None if augmentation_config: @@ -71,12 +51,10 @@ class EncoderDataset(Dataset): if self.verbose: print("\n > DataLoader initialization") print(f" | > Classes per Batch: {num_classes_in_batch}") - print(f" | > Storage Size: {storage_max_size} instances, each with {num_utter_per_class} utters") - print(f" | > Sample_from_storage_p : {self.sample_from_storage_p}") print(f" | > Number of instances : {len(self.items)}") print(f" | > Sequence length: {self.seq_len}") print(f" | > Num Classes: {len(self.classes)}") - print(f" | > Classes: {list(self.classes)}") + print(f" | > Classes: {self.classes}") def load_wav(self, filename): @@ -84,173 +62,84 @@ class EncoderDataset(Dataset): return audio def __parse_items(self): - self.class_to_utters = {} + class_to_utters = {} for i in self.items: path_ = i["audio_file"] speaker_ = i["speaker_name"] - if speaker_ in self.speaker_to_utters.keys(): - self.speaker_to_utters[speaker_].append(path_) + if class_name in class_to_utters.keys(): + class_to_utters[class_name].append(path_) else: - self.class_to_utters[class_name] = [ + class_to_utters[class_name] = [ path_, ] - if self.skip_classes: - self.class_to_utters = { - k: v for (k, v) in self.class_to_utters.items() if len(v) >= self.num_utter_per_class - } + # skip classes with number of samples >= self.num_utter_per_class + class_to_utters = { + k: v for (k, v) in class_to_utters.items() if len(v) >= self.num_utter_per_class + } - self.classes = [k for (k, v) in self.class_to_utters.items()] + self.classes = list(class_to_utters.keys()) + self.classes.sort() + new_items = [] + for item in self.items: + path_ = item[1] + class_name = item[2] + # ignore filtered classes + if class_name not in self.classes: + continue + # ignore small audios + if self.load_wav(path_).shape[0] - self.seq_len <= 0: + continue + + new_items.append({"wav_file_path": path_, "class_name": class_name}) + + self.items = new_items + + def __len__(self): - return int(1e10) + return len(self.items) def get_num_classes(self): return len(self.classes) + def get_class_list(self): + return list(self.classes) + def get_map_classid_to_classname(self): return dict((c_id, c_n) for c_n, c_id in self.classname_to_classid.items()) - def __sample_class(self, ignore_classes=None): - class_name = random.sample(self.classes, 1)[0] - # if list of classes_id is provide make sure that it's will be ignored - if ignore_classes and self.classname_to_classid[class_name] in ignore_classes: - while True: - class_name = random.sample(self.classes, 1)[0] - if self.classname_to_classid[class_name] not in ignore_classes: - break + def __getitem__(self, idx): + return self.items[idx] - if self.num_utter_per_class > len(self.class_to_utters[class_name]): - utters = random.choices(self.class_to_utters[class_name], k=self.num_utter_per_class) - else: - utters = random.sample(self.class_to_utters[class_name], self.num_utter_per_class) - return class_name, utters - - def __sample_class_utterances(self, class_name): - """ - Sample all M utterances for the given class_name. - """ - wavs = [] + def collate_fn(self, batch): + # get the batch class_ids labels = [] - for _ in range(self.num_utter_per_class): - # TODO:dummy but works - while True: - # remove classes that have num_utter less than 2 - if len(self.class_to_utters[class_name]) > 1: - utter = random.sample(self.class_to_utters[class_name], 1)[0] - else: - if class_name in self.classes: - self.classes.remove(class_name) + feats = [] + for item in batch: + utter_path = item["wav_file_path"] + class_name = item["class_name"] - class_name, _ = self.__sample_class() - continue - - wav = self.load_wav(utter) - if wav.shape[0] - self.seq_len > 0: - break - - if utter in self.class_to_utters[class_name]: - self.class_to_utters[class_name].remove(utter) + # get classid + class_id = self.classname_to_classid[class_name] + # load wav file + wav = self.load_wav(utter_path) + offset = random.randint(0, wav.shape[0] - self.seq_len) + wav = wav[offset : offset + self.seq_len] if self.augmentator is not None and self.data_augmentation_p: if random.random() < self.data_augmentation_p: wav = self.augmentator.apply_one(wav) - wavs.append(wav) - labels.append(self.classname_to_classid[class_name]) - return wavs, labels - - def __getitem__(self, idx): - class_name, _ = self.__sample_class() - class_id = self.classname_to_classid[class_name] - return class_name, class_id - - def __load_from_disk_and_storage(self, class_name): - # don't sample from storage, but from HDD - wavs_, labels_ = self.__sample_class_utterances(class_name) - # put the newly loaded item into storage - if self.use_storage: - self.storage.append((wavs_, labels_)) - return wavs_, labels_ - - def collate_fn(self, batch): - # get the batch class_ids - batch = np.array(batch) - classes_id_in_batch = set(batch[:, 1].astype(np.int32)) - - labels = [] - feats = [] - classes = set() - - for class_name, class_id in batch: - class_id = int(class_id) - - # ensure that an class appears only once in the batch - if class_id in classes: - - # remove current class - if class_id in classes_id_in_batch: - classes_id_in_batch.remove(class_id) - - class_name, _ = self.__sample_class(ignore_classes=classes_id_in_batch) - class_id = self.classname_to_classid[class_name] - classes_id_in_batch.add(class_id) - - if self.use_storage and random.random() < self.sample_from_storage_p and self.storage.full(): - # sample from storage (if full) - wavs_, labels_ = self.storage.get_random_sample_fast() - - # force choose the current class or other not in batch - # It's necessary for ideal training with AngleProto and GE2E losses - if labels_[0] in classes_id_in_batch and labels_[0] != class_id: - attempts = 0 - while True: - wavs_, labels_ = self.storage.get_random_sample_fast() - if labels_[0] == class_id or labels_[0] not in classes_id_in_batch: - break - - attempts += 1 - # Try 5 times after that load from disk - if attempts >= 5: - wavs_, labels_ = self.__load_from_disk_and_storage(class_name) - break + if not self.use_torch_spec: + mel = self.ap.melspectrogram(wav) + feats.append(torch.FloatTensor(mel)) else: - # don't sample from storage, but from HDD - wavs_, labels_ = self.__load_from_disk_and_storage(class_name) + feats.append(torch.FloatTensor(wav)) - # append class for control - classes.add(labels_[0]) - - # remove current class and append other - if class_id in classes_id_in_batch: - classes_id_in_batch.remove(class_id) - - classes_id_in_batch.add(labels_[0]) - - # get a random subset of each of the wavs and extract mel spectrograms. - feats_ = [] - for wav in wavs_: - offset = random.randint(0, wav.shape[0] - self.seq_len) - wav = wav[offset : offset + self.seq_len] - # add random gaussian noise - if self.gaussian_augmentation_config and self.gaussian_augmentation_config["p"]: - if random.random() < self.gaussian_augmentation_config["p"]: - wav += np.random.normal( - self.gaussian_augmentation_config["min_amplitude"], - self.gaussian_augmentation_config["max_amplitude"], - size=len(wav), - ) - - if not self.use_torch_spec: - mel = self.ap.melspectrogram(wav) - feats_.append(torch.FloatTensor(mel)) - else: - feats_.append(torch.FloatTensor(wav)) - - labels.append(torch.LongTensor(labels_)) - feats.extend(feats_) + labels.append(class_id) feats = torch.stack(feats) - labels = torch.stack(labels) + labels = torch.LongTensor(labels) return feats, labels diff --git a/TTS/encoder/utils/samplers.py b/TTS/encoder/utils/samplers.py new file mode 100644 index 00000000..d54a5a2e --- /dev/null +++ b/TTS/encoder/utils/samplers.py @@ -0,0 +1,100 @@ +import torch +import random +from torch.utils.data.sampler import Sampler, SubsetRandomSampler + + +class SubsetSampler(Sampler): + """ + Samples elements sequentially from a given list of indices. + + Args: + indices (list): a sequence of indices + """ + + def __init__(self, indices): + self.indices = indices + + def __iter__(self): + return (self.indices[i] for i in range(len(self.indices))) + + def __len__(self): + return len(self.indices) + + +class PerfectBatchSampler(Sampler): + """ + Samples a mini-batch of indices for a balanced class batching + + Args: + dataset_items(list): dataset items to sample from. + classes (list): list of classes of dataset_items to sample from. + batch_size (int): total number of samples to be sampled in a mini-batch. + num_gpus (int): number of GPU in the data parallel mode. + shuffle (bool): if True, samples randomly, otherwise samples sequentially. + drop_last (bool): if True, drops last incomplete batch. + """ + + def __init__(self, dataset_items, classes, batch_size, num_classes_in_batch, num_gpus=1, shuffle=True, drop_last=False): + + assert batch_size % (len(classes) * num_gpus) == 0, ( + 'Batch size must be divisible by number of classes times the number of data parallel devices (if enabled).') + + label_indices = {} + for idx in range(len(dataset_items)): + label = dataset_items[idx]['class_name'] + if label not in label_indices: label_indices[label] = [] + label_indices[label].append(idx) + + if shuffle: + self._samplers = [SubsetRandomSampler(label_indices[key]) for key in classes] + else: + self._samplers = [SubsetSampler(label_indices[key]) for key in classes] + + self._batch_size = batch_size + self._drop_last = drop_last + self._dp_devices = num_gpus + self._num_classes_in_batch = num_classes_in_batch + + def __iter__(self): + + batch = [] + if self._num_classes_in_batch != len(self._samplers): + valid_samplers_idx = random.sample(range(len(self._samplers)), self._num_classes_in_batch) + else: + valid_samplers_idx = None + + iters = [iter(s) for s in self._samplers] + done = False + + while True: + b = [] + for i in range(len(iters)): + if valid_samplers_idx is not None and i not in valid_samplers_idx: + continue + it = iters[i] + idx = next(it, None) + if idx is None: + done = True + break + b.append(idx) + if done: break + batch += b + if len(batch) == self._batch_size: + yield batch + batch = [] + if valid_samplers_idx is not None: + valid_samplers_idx = random.sample(range(len(self._samplers)), self._num_classes_in_batch) + + if not self._drop_last: + if len(batch) > 0: + groups = len(batch) // self._num_classes_in_batch + if groups % self._dp_devices == 0: + yield batch + else: + batch = batch[:(groups // self._dp_devices) * self._dp_devices * self._num_classes_in_batch] + if len(batch) > 0: + yield batch + + def __len__(self): + class_batch_size = self._batch_size // self._num_classes_in_batch + return min(((len(s) + class_batch_size - 1) // class_batch_size) for s in self._samplers) \ No newline at end of file