File size: 4,920 Bytes
167f5d0 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import logging
import torch
from funasr.register import tables
# @tables.register("dataloader_classes", "DataloaderMapStyle")
def DataloaderMapStyle(frontend=None, tokenizer=None, **kwargs):
# dataset
logging.info("Build dataloader")
dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset"))
dataset_tr = dataset_class(
kwargs.get("train_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=True,
**kwargs.get("dataset_conf"),
)
dataset_val = dataset_class(
kwargs.get("valid_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=False,
**kwargs.get("dataset_conf"),
)
# dataloader
batch_sampler = kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
batch_sampler_val = None
if batch_sampler is not None:
batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
batch_sampler = batch_sampler_class(dataset_tr, **kwargs.get("dataset_conf"))
batch_sampler_val = batch_sampler_class(
dataset_val, is_training=False, **kwargs.get("dataset_conf")
)
dataloader_tr = torch.utils.data.DataLoader(
dataset_tr, collate_fn=dataset_tr.collator, **batch_sampler
)
dataloader_val = torch.utils.data.DataLoader(
dataset_val, collate_fn=dataset_val.collator, **batch_sampler_val
)
return dataloader_tr, dataloader_val
@tables.register("dataloader_classes", "DataloaderMapStyle")
class DataloaderMapStyle:
def __init__(self, frontend=None, tokenizer=None, **kwargs):
# dataset
logging.info("Build dataloader")
dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "AudioDataset"))
dataset_tr = None
# split dataset
self.data_split_num = kwargs["dataset_conf"].get("data_split_num", 1)
if self.data_split_num == 1:
dataset_tr = dataset_class(
kwargs.get("train_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=True,
**kwargs.get("dataset_conf"),
)
dataset_val = dataset_class(
kwargs.get("valid_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=False,
**kwargs.get("dataset_conf"),
)
self.dataset_tr = dataset_tr
self.dataset_val = dataset_val
self.kwargs = kwargs
self.dataset_class = dataset_class
self.frontend = frontend
self.tokenizer = tokenizer
self.kwargs = kwargs
def build_iter(self, epoch=0, data_split_i=0, start_step=0, **kwargs):
# reload dataset slice
if self.data_split_num > 1:
del self.dataset_tr
self.dataset_tr = self.dataset_class(
self.kwargs.get("train_data_set_list"),
frontend=self.frontend,
tokenizer=self.tokenizer,
is_training=True,
**self.kwargs.get("dataset_conf"),
data_split_i=data_split_i,
)
# dataloader
batch_sampler = self.kwargs["dataset_conf"].get("batch_sampler", "BatchSampler")
batch_sampler_val = None
if batch_sampler is not None:
batch_sampler_class = tables.batch_sampler_classes.get(batch_sampler)
batch_sampler = batch_sampler_class(
self.dataset_tr, start_step=start_step, **self.kwargs.get("dataset_conf")
)
batch_sampler_val = batch_sampler_class(
self.dataset_val, is_training=False, **self.kwargs.get("dataset_conf")
)
batch_sampler["batch_sampler"].set_epoch(epoch)
batch_sampler_val["batch_sampler"].set_epoch(epoch)
dataloader_tr = torch.utils.data.DataLoader(
self.dataset_tr, collate_fn=self.dataset_tr.collator, **batch_sampler
)
dataloader_val = torch.utils.data.DataLoader(
self.dataset_val, collate_fn=self.dataset_val.collator, **batch_sampler_val
)
return dataloader_tr, dataloader_val
@tables.register("dataloader_classes", "DataloaderIterable")
def DataloaderIterable(frontend=None, tokenizer=None, **kwargs):
logging.info("Build dataloader")
dataset_class = tables.dataset_classes.get(kwargs.get("dataset", "LargeDataset"))
dataset_tr = dataset_class(
kwargs.get("train_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=True,
**kwargs.get("dataset_conf"),
)
dataset_val = dataset_class(
kwargs.get("valid_data_set_list"),
frontend=frontend,
tokenizer=tokenizer,
is_training=False,
**kwargs.get("dataset_conf"),
)
return dataset_tr, dataset_val
|