File size: 16,527 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 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 | """Wrapper around a LightningDataModule."""
import logging
from functools import partial
from typing import Any, Dict, List, Optional
from lightning import LightningDataModule
from lightning.pytorch.trainer.states import TrainerFn
from torch.utils.data import DataLoader
from src.data.loading.components.custom_dataloader import DataloaderWithIterationRetry
from src.data.loading.components.interfaces import BaseDataloaderConfig
from src.data.loading.utils import assign_files_to_workers
from src.utils.file_utils import list_files
class SequenceDataModule(LightningDataModule):
"""A LightningDataModule that encapsulates data splitting, preprocessing,
parallelization and batching.
"""
def __init__(
self,
train_dataloader_config: Optional[BaseDataloaderConfig] = None,
val_dataloader_config: Optional[BaseDataloaderConfig] = None,
test_dataloader_config: Optional[BaseDataloaderConfig] = None,
predict_dataloader_config: Optional[BaseDataloaderConfig] = None,
) -> None:
"""Construct a SequenceDataModule using the provided config files.
The attributes `map_train_files_per_device`,
`map_val_files_per_device`, and `map_test_files_per_device` are
initialized as None, and are later modified by `setup()` to contain
mappings from device indices to lists of data files assigned to that
device.
:param train_dataloader_config: Training dataloader configuration passed
by Hydra.
:param val_dataloader_config: Validation dataloader configuration passed
by Hydra.
:param test_dataloader_config: Test dataloader configuration passed
by Hydra.
"""
super().__init__()
# this line allows to access init params with 'self.hparams' attribute
# also ensures init params will be stored in ckpt
self.save_hyperparameters(logger=False)
self.stage_to_config = {
TrainerFn.FITTING: train_dataloader_config,
TrainerFn.VALIDATING: val_dataloader_config,
TrainerFn.TESTING: test_dataloader_config,
TrainerFn.PREDICTING: predict_dataloader_config,
}
self.stage_to_file_map: Dict[TrainerFn, Dict[int, List[str]]] = dict()
def _get_partial_collate_fn(
self, dataloader_config: BaseDataloaderConfig
) -> callable:
"""Prepare the collate function for the dataloader.
:param config: The dataloader configuration.
"""
# We need to set the collate function here because we can't pickle
# lambda functions but the collate fn needs to receive just the batch.
partial_collate_fn = partial(
dataloader_config.collate_fn,
labels=dataloader_config.labels,
sequence_length=dataloader_config.sequence_length,
masking_token=dataloader_config.masking_token,
padding_token=dataloader_config.padding_token,
oov_token=dataloader_config.get("oov_token", None),
)
return partial_collate_fn
def get_file_suffix_from_config(self, config) -> str:
return (
config.dataset_config.file_format
if getattr(config.dataset_config, "file_format", None)
else config.dataset_config.data_iterator.get_file_suffix()
)
def setup(self, stage: Optional[str] = None) -> None:
"""Assign data files to GPUs.
Note that `self.trainer.world_size` is the total number of GPUs, and is
equal to (# nodes) x (# GPUs per node).
This method is called by Lightning before `trainer.fit()`,
`trainer.validate()`, `trainer.test()`, and `trainer.predict()`, so be
careful not to execute things like random split twice! Also, it is
called after `self.prepare_data()` and there is a barrier in between
which ensures that all the processes proceed to `self.setup()` once the
data is prepared and available for use.
:param stage: Unused parameter. Lightning implementation of setup() uses
stage to determine which dataset splits (train, val, test) to set
up, but we choose to set up all splits on each call to setup() here.
:raise AttributeError: If `self.trainer` is not initialized.
"""
if not hasattr(self, "trainer") or self.trainer is None:
raise AttributeError(
f"self.trainer must be initialized before call to setup()."
)
for stage, config in self.stage_to_config.items():
if (
config is None
): # config is None when we don't want to set up the stages. ie. For inference, we only initialize the predict stage.
self.stage_to_file_map[stage] = {}
else:
# If the stage has not been initialized yet, we assign files to workers based on the suffix passed by the config.
if stage not in self.stage_to_file_map:
list_of_files = list_files(
folder_path=config.data_folder,
suffix=f"*{self.get_file_suffix_from_config(config)}",
)
if hasattr(config, "limit_files") and config.limit_files:
list_of_files = list_of_files[: config.limit_files]
self.stage_to_file_map[stage], _ = assign_files_to_workers(
list_of_files=list_of_files,
total_workers=self.trainer.world_size,
assign_by_size=config.assign_files_by_size,
should_shuffle_rows=config.should_shuffle_rows
if hasattr(config, "should_shuffle_rows")
else False,
assign_all_files_per_worker=config.assign_all_files_per_worker
if hasattr(config, "assign_all_files_per_worker")
else False,
)
def get_dataloader(self, stage: TrainerFn) -> DataLoader[Any]:
"""Construct a DataLoader on a single GPU using config `curr_config`.
The single GPU is managed by Lightning and corresponds to
`self.trainer.global_rank`.
:param curr_config: Config that determines dataloader properties.
:param map_files_idx_per_device: Map from GPUs to file indices.
:return: DataLoader running on one GPU that processes the files
assigned to that GPU, according to `map_files_per_device`.
:raise AttributeError: If `self.trainer` is not initialized or if setup was not called.
"""
if not hasattr(self, "trainer"):
raise AttributeError(
f"self.trainer must be initialized before call to get_dataloader()."
)
if not self.stage_to_file_map[stage]:
raise AttributeError(f"Stage {stage} must initialize file map.")
curr_config = self.stage_to_config[stage]
assign_all_files_per_worker = (
curr_config.assign_all_files_per_worker
if hasattr(curr_config, "assign_all_files_per_worker")
else False
)
# We initialize the dataset with the parameters passed on the config.
dataset = curr_config.dataset_class(
dataset_config=curr_config.dataset_config,
data_folder=curr_config.data_folder,
should_shuffle_rows=curr_config.should_shuffle_rows,
batch_size=curr_config.batch_size_per_device,
is_for_training=stage == TrainerFn.FITTING,
assign_all_files_per_worker=assign_all_files_per_worker,
) # type: ignore
device_file_list = self.stage_to_file_map[stage].get(
self.trainer.global_rank, []
)
dataset.set_list_of_files(list_of_files=device_file_list)
# set the number of total GPUs and GPU index
dataset.set_distributed_params(
total_workers=self.trainer.world_size,
global_worker_id=self.trainer.global_rank,
)
# Any additional parameters for the masking function should be added to
# the config and passed there. This is required because we can't pickle
# lambda functions but the collate fn needs to receive just the batch.
collate_fn_partial = self._get_partial_collate_fn(curr_config)
if curr_config.num_workers == 0:
persistent_workers = False
logging.warning(
"num_workers is set to 0, persistent_workers will be set to"
" False as persistent workers require num_workers > 0"
)
else:
persistent_workers = curr_config.persistent_workers
return (
DataloaderWithIterationRetry(
dataset=dataset,
batch_size=curr_config.batch_size_per_device
if curr_config.dataset_config.iterate_per_row
else None,
num_workers=curr_config.num_workers, # num workers per GPU
pin_memory=curr_config.pin_memory,
persistent_workers=persistent_workers,
drop_last=curr_config.drop_last
if curr_config.dataset_config.iterate_per_row
else False,
collate_fn=collate_fn_partial,
timeout=curr_config.timeout,
),
) # type: ignore
def train_dataloader(self) -> DataLoader[Any]:
"""Create and return the train dataloader.
:return: The train dataloader.
"""
return self.get_dataloader(stage=TrainerFn.FITTING)
def val_dataloader(self) -> DataLoader[Any]:
"""Create and return the validation dataloader.
:return: The validation dataloader.
"""
return self.get_dataloader(stage=TrainerFn.VALIDATING)
def test_dataloader(self) -> DataLoader[Any]:
"""Create and return the test dataloader.
:return: The test dataloader.
"""
return self.get_dataloader(stage=TrainerFn.TESTING)
def predict_dataloader(self) -> DataLoader[Any]:
"""Create and return the predict dataloader.
:return: The predict dataloader.
"""
return self.get_dataloader(stage=TrainerFn.PREDICTING)
def teardown(self, stage: Optional[str] = None) -> None:
"""Lightning hook for cleaning up after `trainer.fit()`,
`trainer.validate()`, `trainer.test()`, and `trainer.predict()`.
:param stage: The stage being torn down. Defaults to ``None``.
"""
pass
def state_dict(self) -> Dict[Any, Any]:
"""Called when saving a checkpoint. Implement to generate and save the
datamodule state.
:return: A dictionary containing the datamodule state that you want to
save.
"""
return {}
def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
"""Called when loading a checkpoint. Implement to reload datamodule
state given datamodule `state_dict()`.
:param state_dict: The datamodule state returned by `self.state_dict()`.
"""
pass
class ItemDataModule(SequenceDataModule):
"""A LightningDataModule that encapsulates data splitting, preprocessing,
parallelization and batching individual (non-sequential) item features.
"""
def __init__(
self,
train_dataloader_config: Optional[BaseDataloaderConfig] = None,
val_dataloader_config: Optional[BaseDataloaderConfig] = None,
test_dataloader_config: Optional[BaseDataloaderConfig] = None,
predict_dataloader_config: Optional[BaseDataloaderConfig] = None,
) -> None:
"""Construct a SequenceDataModule using the provided config files.
The attributes `map_train_files_per_device`,
`map_val_files_per_device`, and `map_test_files_per_device` are
initialized as None, and are later modified by `setup()` to contain
mappings from device indices to lists of data files assigned to that
device.
:param train_dataloader_config: Training dataloader configuration passed
by Hydra.
:param val_dataloader_config: Validation dataloader configuration passed
by Hydra.
:param test_dataloader_config: Test dataloader configuration passed
by Hydra.
:param predict_dataloader_config: Prediction dataloader configuration
passed by Hydra.
"""
super().__init__(
train_dataloader_config=train_dataloader_config,
val_dataloader_config=val_dataloader_config,
test_dataloader_config=test_dataloader_config,
predict_dataloader_config=predict_dataloader_config,
)
def get_dataloader(self, stage: TrainerFn) -> DataLoader[Any]:
"""Construct a DataLoader on a single GPU using config `curr_config`.
The single GPU is managed by Lightning and corresponds to
`self.trainer.global_rank`.
:param curr_config: Config that determines dataloader properties.
:param map_files_idx_per_device: Map from GPUs to file indices.
:return: DataLoader running on one GPU that processes the files
assigned to that GPU, according to `map_files_per_device`.
:raise AttributeError: If `self.trainer` is not initialized or if setup was not called.
"""
if not hasattr(self, "trainer"):
raise AttributeError(
f"self.trainer must be initialized before call to get_dataloader()."
)
if not self.stage_to_file_map[stage]:
raise AttributeError(f"Stage {stage} must initialize file map.")
curr_config = self.stage_to_config[stage]
assign_all_files_per_worker = (
curr_config.assign_all_files_per_worker
if hasattr(curr_config, "assign_all_files_per_worker")
else False
)
assert (
stage == TrainerFn.FITTING or not assign_all_files_per_worker
), "Automatic file assignment should only be disabled for training."
# We initialize the dataset with the parameters passed on the config.
dataset = curr_config.dataset_class(
dataset_config=curr_config.dataset_config,
data_folder=curr_config.data_folder,
should_shuffle_rows=curr_config.should_shuffle_rows,
batch_size=curr_config.batch_size_per_device,
is_for_training=stage == TrainerFn.FITTING,
assign_all_files_per_worker=assign_all_files_per_worker,
) # type: ignore
device_file_list = self.stage_to_file_map[stage].get(
self.trainer.global_rank, []
)
dataset.set_list_of_files(list_of_files=device_file_list)
# set the number of total GPUs and GPU index
dataset.set_distributed_params(
total_workers=self.trainer.world_size,
global_worker_id=self.trainer.global_rank,
)
if curr_config.num_workers == 0:
persistent_workers = False
logging.warning(
"num_workers is set to 0, persistent_workers will be set to"
" False as persistent workers require num_workers > 0"
)
else:
persistent_workers = curr_config.persistent_workers
return (
DataloaderWithIterationRetry(
dataset=dataset,
batch_size=curr_config.batch_size_per_device
if curr_config.dataset_config.iterate_per_row
else None,
num_workers=curr_config.num_workers, # num workers per GPU
pin_memory=curr_config.pin_memory,
persistent_workers=persistent_workers,
drop_last=curr_config.drop_last
if curr_config.dataset_config.iterate_per_row
else False,
collate_fn=curr_config.collate_fn,
timeout=curr_config.timeout,
),
) # type: ignore
|