From 968e740629259a863ee6af9b385121b9f7e51c64 Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:43:44 +0200 Subject: [PATCH 1/2] Add lightning batch size finder --- configs/callbacks/batch_size_finder.yaml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 configs/callbacks/batch_size_finder.yaml diff --git a/configs/callbacks/batch_size_finder.yaml b/configs/callbacks/batch_size_finder.yaml new file mode 100644 index 0000000..55533d3 --- /dev/null +++ b/configs/callbacks/batch_size_finder.yaml @@ -0,0 +1,5 @@ +# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.BatchSizeFinder.html + +_target_: lightning.pytorch.callbacks.BatchSizeFinder +mode: 'binsearch' +init_val: 16 \ No newline at end of file From 117082d7a45b58ec4b1e0c4c38adb1173031a125 Mon Sep 17 00:00:00 2001 From: GabrieleTi Date: Wed, 20 May 2026 09:44:16 +0200 Subject: [PATCH 2/2] Implement dynamic batch size calculation and add pre-fetch param --- src/data/base_datamodule.py | 34 +++++++++++++++++++++++----------- 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/src/data/base_datamodule.py b/src/data/base_datamodule.py index 6d1147a..437de8a 100644 --- a/src/data/base_datamodule.py +++ b/src/data/base_datamodule.py @@ -27,6 +27,8 @@ def __init__( train_val_test_split: Tuple[float, float, float] = (0.7, 0.15, 0.15), num_workers: int = 0, pin_memory: bool = False, + persistent_workers: bool = False, + prefetch_factor: int | None = None, dataset_name: str = "base", split_mode: str = "random", save_split: bool = False, @@ -43,6 +45,8 @@ def __init__( testing :param num_workers: number of workers for dataloader :param pin_memory: pin memory for dataloader + :param persistent_workers: keep DataLoader workers alive between epochs + :param persistent_workers: :param dataset_name: dataset name :param split_mode: data split mode: random/from_file :param save_split: if to save split file @@ -55,7 +59,6 @@ def __init__( self.save_hyperparameters(logger=False) self.dataset: BaseDataset = dataset - self.batch_size_per_device: int = batch_size # Caption generation self.use_collate_fn: bool = self.dataset.use_aux_data @@ -84,16 +87,19 @@ def setup(self, stage: str = "fit") -> None: # Set up the dataset (download requested modalities) self.dataset.setup() - self.setup_batch_size_per_device() - def setup_batch_size_per_device(self) -> None: + @property + def batch_size_per_device(self) -> None: """Divide batch size by the number of devices.""" - if self.trainer is not None: - if self.hparams.batch_size % self.trainer.world_size != 0: - raise RuntimeError( - f"Batch size ({self.hparams.batch_size}) is not divisible by the number of devices ({self.trainer.world_size})." - ) - self.batch_size_per_device = self.hparams.batch_size // self.trainer.world_size + if self.trainer is None: + return self.hparams.batch_size + + if self.hparams.batch_size % self.trainer.world_size != 0: + raise RuntimeError(f"Batch size ({self.hparams.batch_size}) is not divisible by the number of devices ({self.trainer.world_size}).") + + return self.hparams.batch_size // self.trainer.world_size + + def split_data(self) -> None: """Split data into train, val and test. @@ -355,9 +361,12 @@ def train_dataloader(self) -> DataLoader[Any]: return DataLoader( dataset=self.data_train, batch_size=self.batch_size_per_device, - persistent_workers=True if self.hparams.num_workers > 0 else False, + persistent_workers=( + bool(self.hparams.persistent_workers) if self.hparams.num_workers > 0 else False + ), num_workers=self.hparams.num_workers, pin_memory=self.hparams.pin_memory, + prefetch_factor=(self.hparams.prefetch_factor if self.hparams.num_workers > 0 else None), shuffle=True, collate_fn=( partial( @@ -380,7 +389,8 @@ def val_dataloader(self) -> DataLoader[Any]: batch_size=self.batch_size_per_device, num_workers=self.hparams.num_workers, pin_memory=self.hparams.pin_memory, - persistent_workers=True if self.hparams.num_workers > 0 else False, + persistent_workers=(bool(self.hparams.persistent_workers) if self.hparams.num_workers > 0 else False), + prefetch_factor=(self.hparams.prefetch_factor if self.hparams.num_workers > 0 else None), shuffle=False, collate_fn=( partial( @@ -403,6 +413,8 @@ def test_dataloader(self) -> DataLoader[Any]: batch_size=self.batch_size_per_device, num_workers=self.hparams.num_workers, pin_memory=self.hparams.pin_memory, + persistent_workers=(bool(self.hparams.persistent_workers) if self.hparams.num_workers > 0 else False), + prefetch_factor=(self.hparams.prefetch_factor if self.hparams.num_workers > 0 else None), shuffle=False, collate_fn=( partial(