File size: 1,267 Bytes
32b2aaa |
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 |
import logging
import random
from torch.utils.data import DataLoader
from ..hparams import HParams
from .dataset import Dataset
from .utils import mix_fg_bg, rglob_audio_files
logger = logging.getLogger(__name__)
def _create_datasets(hp: HParams, mode, val_size=10, seed=123):
paths = rglob_audio_files(hp.fg_dir)
logger.info(f"Found {len(paths)} audio files in {hp.fg_dir}")
random.Random(seed).shuffle(paths)
train_paths = paths[:-val_size]
val_paths = paths[-val_size:]
train_ds = Dataset(train_paths, hp, training=True, mode=mode)
val_ds = Dataset(val_paths, hp, training=False, mode=mode)
logger.info(f"Train set: {len(train_ds)} samples - Val set: {len(val_ds)} samples")
return train_ds, val_ds
def create_dataloaders(hp: HParams, mode):
train_ds, val_ds = _create_datasets(hp=hp, mode=mode)
train_dl = DataLoader(
train_ds,
batch_size=hp.batch_size_per_gpu,
shuffle=True,
num_workers=hp.nj,
drop_last=True,
collate_fn=train_ds.collate_fn,
)
val_dl = DataLoader(
val_ds,
batch_size=1,
shuffle=False,
num_workers=hp.nj,
drop_last=False,
collate_fn=val_ds.collate_fn,
)
return train_dl, val_dl
|