File size: 7,356 Bytes
d9115cb | 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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | import copy
from typing import Any, List, Optional
from torch.utils.data import IterableDataset, get_worker_info
from src.data.loading.components.interfaces import BaseDatasetConfig
from src.utils.pylogger import RankedLogger
command_line_logger = RankedLogger(__name__, rank_zero_only=True)
class BaseDataset:
def __init__(
self,
dataset_config: BaseDatasetConfig,
data_folder: str,
should_shuffle_rows: bool = False,
batch_size: int = 1,
is_for_training: bool = True,
assign_all_files_per_worker: bool = False,
):
"""
Base class for all datasets. This class is used to set up the dataset and provide the list of files to be used.
Args:
dataset_config (BaseDatasetConfig): Configuration for the dataset.
data_folder (str): Path to the folder where the data is stored.
should_shuffle_rows (bool): Whether to shuffle the rows of the dataset.
batch_size (int): Batch size to be used for the dataset.
is_for_training (bool): Whether the dataset is for training or not.
assign_all_files_per_worker (bool): Whether to assign all files to each worker or not.
This will enable each worker to access all files. Each worker will locally shuffle the files.
This would be useful for small datasets. In smaller datasets, if each worker only observes a subset of the files,
it may not be able to learn the distribution of the data.
"""
self.dataset_config = dataset_config
self.should_shuffle_rows = should_shuffle_rows
self.data_folder = data_folder
self.list_of_file_paths = []
self.batch_size = batch_size
self.is_for_training = is_for_training
self.assign_all_files_per_worker = assign_all_files_per_worker
def set_list_of_files(self, list_of_files: List[str]):
self.list_of_file_paths = list_of_files
def set_distributed_params(self, total_workers: int, global_worker_id: int):
# TODO (lneves): figure out how to do this in LightningDataModule
self.total_workers = total_workers
self.global_worker_id = global_worker_id
def get_worker_id_and_num_workers(self):
worker_info = get_worker_info()
if worker_info is None:
# Single-worker setup (no multiprocessing)
worker_id = 0
num_workers = 1
else:
# Multi-worker setup
worker_id = worker_info.id
num_workers = worker_info.num_workers
self.global_dataloader_worker_id = (
self.global_worker_id * num_workers + worker_id
)
return worker_id, num_workers
def get_list_of_worker_files(self):
# Get information about worker and then separate only files that belong to this worker
worker_id, num_workers = self.get_worker_id_and_num_workers()
if self.assign_all_files_per_worker:
worker_files = self.list_of_file_paths
else:
worker_files = self.list_of_file_paths[worker_id::num_workers]
command_line_logger.debug(
f"GPU Worker: {self.global_worker_id}/{self.total_workers} CPU Worker {worker_id} has {len(worker_files)} files"
)
return worker_files
def setup(self):
pass
class UnboundedSequenceIterable(BaseDataset, IterableDataset):
"""An unbounded dataset is a dataset that we don't know the size of beforehand.
For training, we will iterate over the dataset infinitely. For evaluation, we will iterate over the dataset once.
"""
def __init__(
self,
dataset_config: BaseDatasetConfig,
data_folder: str,
should_shuffle_rows: bool = False,
batch_size: int = 1,
is_for_training: bool = True,
assign_all_files_per_worker: bool = False,
):
super().__init__(
dataset_config=dataset_config,
data_folder=data_folder,
should_shuffle_rows=should_shuffle_rows,
batch_size=batch_size,
is_for_training=is_for_training,
assign_all_files_per_worker=assign_all_files_per_worker,
)
self.data_iterator = dataset_config.data_iterator
self.dataset_to_iterate = None
def setup(self):
# We update each worker's data iterator with the files just for that worker.
self.data_iterator.update_list_of_file_paths(self.get_list_of_worker_files())
self.data_iterator = (
# here we use global_dataloader_worker_id as the seed for shuffling
# this doesn't matter for the case where workers have non-overlapping files
# but it does matter for the case where workers have all files
# (e.g. when using assign_all_files_per_worker)
# the same seed is used for all workers would cause duplicated examples returned by different workers
self.data_iterator.shuffle(seed=self.global_dataloader_worker_id)
if self.should_shuffle_rows
else self.data_iterator
)
self.data_iterator.should_shuffle_rows = self.should_shuffle_rows
# We provide the flexibility to iterate per row, if per row preprocessing is needed, or per batch.
self.dataset_to_iterate = (
self.data_iterator.iterrows()
if self.dataset_config.iterate_per_row
else self.data_iterator.iter_batches(self.batch_size)
)
command_line_logger.debug(
f"GLOBAL ID {self.global_dataloader_worker_id} GPU Worker: {self.global_worker_id}/{self.total_workers} with {len(self.data_iterator.list_of_file_paths)} files\
First five files are: {self.data_iterator.list_of_file_paths[:5]}"
)
def __iter__(self):
if self.dataset_to_iterate is None:
# If it has not been set up, it means it is a forkserver worker. We need to set it up.
self.setup()
# If the dataset is for training, we want to keep iterating over the dataset infinitely.
# On a streaming dataset, we will always be on Epoch 0.
finished_iteration = False
while not finished_iteration:
for row_or_batch in self.dataset_to_iterate:
for (
preprocessing_function
) in self.dataset_config.preprocessing_functions:
row_or_batch = preprocessing_function(
row_or_batch, dataset_config=self.dataset_config
)
if row_or_batch is None:
break
if row_or_batch:
yield row_or_batch
# if the dataset is not for training, we stop the loop. Otherwise, we continue.
finished_iteration = not self.is_for_training
if not finished_iteration:
self.setup()
# We reset the dataset to iterate to None, so that it is set up again in the next iteration.
# This is required for validation when persitent_workers = True.
self.dataset_to_iterate = None
return None |