import os import sys import math import numpy as np import torch from torch.utils.data.sampler import Sampler from torch.utils.data.distributed import DistributedSampler import torch.distributed from typing import TypeVar, Optional, Iterator,List import logging import pandas as pd import glob import torch.distributed as dist logger = logging.getLogger(f'main.{__name__}') sys.path.insert(0, '.') # nopep8 class JoinManifestSpecs(torch.utils.data.Dataset): def __init__(self, split, spec_dir_path, mel_num=80,spec_crop_len=1248,mode='pad',pad_value=-5,drop=0,**kwargs): super().__init__() self.split = split self.max_batch_len = spec_crop_len self.min_batch_len = 64 self.mel_num = mel_num self.min_factor = 4 self.drop = drop self.pad_value = pad_value assert mode in ['pad','tile'] self.collate_mode = mode # print(f"################# self.collate_mode {self.collate_mode} ##################") manifest_files = [] for dir_path in spec_dir_path.split(','): manifest_files += glob.glob(f'{dir_path}/*.tsv') df_list = [pd.read_csv(manifest,sep='\t') for manifest in manifest_files] df = pd.concat(df_list,ignore_index=True) if split == 'train': self.dataset = df.iloc[100:] elif split == 'valid' or split == 'val': self.dataset = df.iloc[:100] elif split == 'test': df = self.add_name_num(df) self.dataset = df else: raise ValueError(f'Unknown split {split}') self.dataset.reset_index(inplace=True) print('dataset len:', len(self.dataset)) def add_name_num(self,df): """each file may have different caption, we add num to filename to identify each audio-caption pair""" name_count_dict = {} change = [] for t in df.itertuples(): name = getattr(t,'name') if name in name_count_dict: name_count_dict[name] += 1 else: name_count_dict[name] = 0 change.append((t[0],name_count_dict[name])) for t in change: df.loc[t[0],'name'] = df.loc[t[0],'name'] + f'_{t[1]}' return df def ordered_indices(self): index2dur = self.dataset[['duration']] index2dur = index2dur.sort_values(by='duration') return list(index2dur.index) def __getitem__(self, idx): item = {} data = self.dataset.iloc[idx] try: spec = np.load(data['mel_path']) # mel spec [80, 624] except: mel_path = data['mel_path'] print(f'corrupted:{mel_path}') spec = np.ones((self.mel_num,self.min_batch_len)).astype(np.float32)*self.pad_value item['image'] = spec p = np.random.uniform(0,1) if p > self.drop: ori_caption = data['caption'] struct_caption = f'<{ori_caption}& all>' else: ori_caption = "" struct_caption = "" item["caption"] = {"ori_caption":ori_caption,"struct_caption":struct_caption} if self.split == 'test': item['f_name'] = data['name'] # item['f_name'] = data['mel_path'] return item def collater(self,inputs): to_dict = {} for l in inputs: for k,v in l.items(): if k in to_dict: to_dict[k].append(v) else: to_dict[k] = [v] if self.collate_mode == 'pad': to_dict['image'] = collate_1d_or_2d(to_dict['image'],pad_idx=self.pad_value,min_len = self.min_batch_len,max_len=self.max_batch_len,min_factor=self.min_factor) elif self.collate_mode == 'tile': to_dict['image'] = collate_1d_or_2d_tile(to_dict['image'],min_len = self.min_batch_len,max_len=self.max_batch_len,min_factor=self.min_factor) else: raise NotImplementedError to_dict['caption'] = {'ori_caption':[c['ori_caption'] for c in to_dict['caption']], 'struct_caption':[c['struct_caption'] for c in to_dict['caption']]} return to_dict def __len__(self): return len(self.dataset) class JoinSpecsTrain(JoinManifestSpecs): def __init__(self, specs_dataset_cfg): super().__init__('train', **specs_dataset_cfg) class JoinSpecsValidation(JoinManifestSpecs): def __init__(self, specs_dataset_cfg): super().__init__('valid', **specs_dataset_cfg) class JoinSpecsTest(JoinManifestSpecs): def __init__(self, specs_dataset_cfg): super().__init__('test', **specs_dataset_cfg) class JoinSpecsDebug(JoinManifestSpecs): def __init__(self, specs_dataset_cfg): super().__init__('valid', **specs_dataset_cfg) self.dataset = self.dataset.iloc[:37] class DDPIndexBatchSampler(Sampler):# 让长度相似的音频的indices合到一个batch中以避免过长的pad def __init__(self, indices ,batch_size, num_replicas: Optional[int] = None, rank: Optional[int] = None, shuffle: bool = True, seed: int = 0, drop_last: bool = False) -> None: if num_replicas is None: if not dist.is_initialized(): # raise RuntimeError("Requires distributed package to be available") print("Not in distributed mode") num_replicas = 1 else: num_replicas = dist.get_world_size() if rank is None: if not dist.is_initialized(): # raise RuntimeError("Requires distributed package to be available") rank = 0 else: rank = dist.get_rank() if rank >= num_replicas or rank < 0: raise ValueError( "Invalid rank {}, rank should be in the interval" " [0, {}]".format(rank, num_replicas - 1)) self.indices = indices self.num_replicas = num_replicas self.rank = rank self.epoch = 0 self.drop_last = drop_last self.batch_size = batch_size self.batches = self.build_batches() print(f"rank: {self.rank}, batches_num {len(self.batches)}") # If the dataset length is evenly divisible by replicas, then there # is no need to drop any data, since the dataset will be split equally. if self.drop_last and len(self.batches) % self.num_replicas != 0: self.batches = self.batches[:len(self.batches)//self.num_replicas*self.num_replicas] if len(self.batches) > self.num_replicas: self.batches = self.batches[self.rank::self.num_replicas] else: # may happen in sanity checking self.batches = [self.batches[0]] print(f"after split batches_num {len(self.batches)}") self.shuffle = shuffle if self.shuffle: self.batches = np.random.permutation(self.batches) self.seed = seed def set_epoch(self,epoch): self.epoch = epoch if self.shuffle: np.random.seed(self.seed+self.epoch) self.batches = np.random.permutation(self.batches) def build_batches(self): batches,batch = [],[] for index in self.indices: batch.append(index) if len(batch) == self.batch_size: batches.append(batch) batch = [] if not self.drop_last and len(batch) > 0: batches.append(batch) return batches def __iter__(self) -> Iterator[List[int]]: for batch in self.batches: yield batch def __len__(self) -> int: return len(self.batches) def set_epoch(self, epoch: int) -> None: r""" Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas use a different random ordering for each epoch. Otherwise, the next iteration of this sampler will yield the same ordering. Args: epoch (int): Epoch number. """ self.epoch = epoch def collate_1d_or_2d(values, pad_idx=0, left_pad=False, shift_right=False,min_len = None, max_len=None,min_factor=None, shift_id=1): if len(values[0].shape) == 1: return collate_1d(values, pad_idx, left_pad, shift_right,min_len, max_len,min_factor, shift_id) else: return collate_2d(values, pad_idx, left_pad, shift_right,min_len,max_len,min_factor) def collate_1d(values, pad_idx=0, left_pad=False, shift_right=False,min_len=None, max_len=None,min_factor=None, shift_id=1): """Convert a list of 1d tensors into a padded 2d tensor.""" size = max(v.size(0) for v in values) if max_len: size = min(size,max_len) if min_len: size = max(size,min_len) if min_factor and (size % min_factor!=0):# size must be the multiple of min_factor size += (min_factor - size % min_factor) res = values[0].new(len(values), size).fill_(pad_idx) def copy_tensor(src, dst): assert dst.numel() == src.numel(), f"dst shape:{dst.shape} src shape:{src.shape}" if shift_right: dst[1:] = src[:-1] dst[0] = shift_id else: dst.copy_(src) for i, v in enumerate(values): copy_tensor(v, res[i][size - len(v):] if left_pad else res[i][:len(v)]) return res def collate_2d(values, pad_idx=0, left_pad=False, shift_right=False, min_len=None,max_len=None,min_factor=None): """Collate 2d for melspec,Convert a list of 2d tensors into a padded 3d tensor,pad in mel_length dimension. values[0] shape: (melbins,mel_length) """ size = max(v.shape[1] for v in values) # if max_len is None else max_len if max_len: size = min(size,max_len) if min_len: size = max(size,min_len) if min_factor and (size % min_factor!=0):# size must be the multiple of min_factor size += (min_factor - size % min_factor) if isinstance(values,np.ndarray): values = torch.FloatTensor(values) if isinstance(values,list): values = [torch.FloatTensor(v) for v in values] res = torch.ones(len(values), values[0].shape[0],size).to(dtype=torch.float32)*pad_idx def copy_tensor(src, dst): assert dst.numel() == src.numel(), f"dst shape:{dst.shape} src shape:{src.shape}" if shift_right: dst[1:] = src[:-1] else: dst.copy_(src) for i, v in enumerate(values): copy_tensor(v[:,:size], res[i][:,size - v.shape[1]:] if left_pad else res[i][:,:v.shape[1]]) return res def collate_1d_or_2d_tile(values, shift_right=False,min_len = None, max_len=None,min_factor=None, shift_id=1): if len(values[0].shape) == 1: return collate_1d_tile(values, shift_right,min_len, max_len,min_factor, shift_id) else: return collate_2d_tile(values, shift_right,min_len,max_len,min_factor) def collate_1d_tile(values, shift_right=False,min_len=None, max_len=None,min_factor=None,shift_id=1): """Convert a list of 1d tensors into a padded 2d tensor.""" size = max(v.size(0) for v in values) if max_len: size = min(size,max_len) if min_len: size = max(size,min_len) if min_factor and (size%min_factor!=0):# size must be the multiple of min_factor size += (min_factor - size % min_factor) res = values[0].new(len(values), size) def copy_tensor(src, dst): assert dst.numel() == src.numel(), f"dst shape:{dst.shape} src shape:{src.shape}" if shift_right: dst[1:] = src[:-1] dst[0] = shift_id else: dst.copy_(src) for i, v in enumerate(values): n_repeat = math.ceil((size + 1) / v.shape[0]) v = torch.tile(v,dims=(1,n_repeat))[:size] copy_tensor(v, res[i]) return res def collate_2d_tile(values, shift_right=False, min_len=None,max_len=None,min_factor=None): """Collate 2d for melspec,Convert a list of 2d tensors into a padded 3d tensor,pad in mel_length dimension. """ size = max(v.shape[1] for v in values) # if max_len is None else max_len if max_len: size = min(size,max_len) if min_len: size = max(size,min_len) if min_factor and (size % min_factor!=0):# size must be the multiple of min_factor size += (min_factor - size % min_factor) if isinstance(values,np.ndarray): values = torch.FloatTensor(values) if isinstance(values,list): values = [torch.FloatTensor(v) for v in values] res = torch.zeros(len(values), values[0].shape[0],size).to(dtype=torch.float32) def copy_tensor(src, dst): assert dst.numel() == src.numel() if shift_right: dst[1:] = src[:-1] else: dst.copy_(src) for i, v in enumerate(values): n_repeat = math.ceil((size + 1) / v.shape[1]) v = torch.tile(v,dims=(1,n_repeat))[:,:size] copy_tensor(v, res[i]) return res