import logging
|
|
import kaldiio
|
import numpy as np
|
import torch
|
from torch.utils.data import DataLoader
|
from torch.utils.data import Dataset
|
|
|
def custom_collate(batch):
|
keys, speech, speaker_labels, orders = zip(*batch)
|
speech = [torch.from_numpy(np.copy(sph)).to(torch.float32) for sph in speech]
|
speaker_labels = [torch.from_numpy(np.copy(spk)).to(torch.float32) for spk in speaker_labels]
|
orders = [torch.from_numpy(np.copy(o)).to(torch.int64) for o in orders]
|
batch = dict(speech=speech, speaker_labels=speaker_labels, orders=orders)
|
|
return keys, batch
|
|
|
class EENDOLADataset(Dataset):
|
def __init__(
|
self,
|
data_file,
|
):
|
self.data_file = data_file
|
with open(data_file) as f:
|
lines = f.readlines()
|
self.samples = [line.strip().split() for line in lines]
|
logging.info("total samples: {}".format(len(self.samples)))
|
|
def __len__(self):
|
return len(self.samples)
|
|
def __getitem__(self, idx):
|
key, speech_path, speaker_label_path = self.samples[idx]
|
speech = kaldiio.load_mat(speech_path)
|
speaker_label = kaldiio.load_mat(speaker_label_path).reshape(speech.shape[0], -1)
|
|
order = np.arange(speech.shape[0])
|
np.random.shuffle(order)
|
|
return key, speech, speaker_label, order
|
|
|
class EENDOLADataLoader:
|
def __init__(self, data_file, batch_size, shuffle=True, num_workers=8):
|
dataset = EENDOLADataset(data_file)
|
self.data_loader = DataLoader(
|
dataset,
|
batch_size=batch_size,
|
collate_fn=custom_collate,
|
shuffle=shuffle,
|
num_workers=num_workers,
|
)
|
|
def build_iter(self, epoch):
|
return self.data_loader
|
