Add batch size finder

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

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(