Spaces:
Runtime error
Runtime error
File size: 3,646 Bytes
96ee597 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
"""Data module.
Copyright PolyAI Limited.
"""
import typing
from pathlib import Path
from typing import List
import lightning.pytorch as pl
from torch.utils import data
from data.collation import GlobalCollater
from data.sampler import RandomBucketSampler
from data.single_speaker_dataset import QuantizeDataset
from utils import breakpoint_on_error
class ConcatDataset(data.ConcatDataset):
def __init__(self, datasets) -> None:
super().__init__(datasets)
self.lengths = []
for dataset in datasets:
self.lengths.extend(dataset.lengths)
class DataModule(pl.LightningDataModule):
def __init__(
self, hp, metapath: List[str], val_metapath: List[str],
world_size, local_rank
):
super().__init__()
self.hp = hp
self.metapath = metapath
self.val_metapath = val_metapath
self.world_size = world_size
self.local_rank = local_rank
self.collater = GlobalCollater(
self.hp.n_codes, self.hp.n_semantic_codes)
def setup(self, stage: str) -> None:
if stage == "fit":
self.train_data = self.concatenate_datasets(
self.metapath, dataset_class=QuantizeDataset
)
if stage == "valid":
self.val_data = []
self.val_data_keys = []
self.prepare_val_datasets()
assert len(self.val_data) > 0
assert len(self.val_data_keys) > 0
@breakpoint_on_error
def concatenate_datasets(
self, metapaths, dataset_class: typing.Type[QuantizeDataset]):
data = []
for _, metapath in enumerate(metapaths):
metapath = Path(metapath)
# assumption that audios and audios-embeddings
# are in the same folder as metapath
datadir = metapath.with_name("audios")
assert datadir.exists()
data.append(
dataset_class(
self.hp,
metapath,
datadir=datadir,
speaker_embedding_dir=None,
)
)
return ConcatDataset(data)
def prepare_val_datasets(self):
for manifest in self.val_metapath:
self.val_data.append(
self.concatenate_datasets(
[manifest], dataset_class=QuantizeDataset)
)
name = Path(manifest).parent.name
self.val_data_keys.append(name)
assert len(self.val_data) == len(self.val_data_keys)
def train_dataloader(self):
length = self.train_data.lengths
sampler = RandomBucketSampler(
self.hp.train_bucket_size,
length,
self.hp.batch_size,
drop_last=True,
distributed=self.hp.distributed,
world_size=self.world_size,
rank=self.local_rank,
)
dataloader = data.DataLoader(
self.train_data,
num_workers=self.hp.nworkers,
batch_sampler=sampler,
collate_fn=self.collater.collate,
pin_memory=True
)
return dataloader
def val_dataloader(self):
val_loaders = []
for dataset in self.val_data:
val_loaders.append(
data.DataLoader(
dataset,
num_workers=self.hp.nworkers,
batch_size=int(self.hp.batch_size),
collate_fn=self.collater.collate,
shuffle=False,
pin_memory=True
)
)
return val_loaders
|