Add configurable pin_memory and prefetch_factor for DataLoader

dingo/core/utils/torchutils.py

@@ -9,9 +9,9 @@
 
 def fix_random_seeds(_):
     """Utility function to set random seeds when using multiple workers for DataLoader."""
-    np.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1))
+    np.random.seed(int(torch.initial_seed()) % (2**32 - 1))
     try:
-        bilby.core.utils.random.seed(int(torch.initial_seed()) % (2 ** 32 - 1))
+        bilby.core.utils.random.seed(int(torch.initial_seed()) % (2**32 - 1))
     except AttributeError:  # In case using an old version of Bilby.
         pass
 
@@ -190,7 +190,9 @@ def build_train_and_test_loaders(
     dataset: torch.utils.data.Dataset,
     train_fraction: float,
     batch_size: int,
-    num_workers: int,
+    num_workers: int = 0,
+    pin_memory: bool = True,
+    prefetch_factor: int = 1,
 ):
     """
     Split the dataset into train and test sets, and build corresponding DataLoaders.
@@ -204,6 +206,14 @@ def build_train_and_test_loaders(
         Should lie between 0 and 1.
     batch_size : int
     num_workers : int
+        Number of workers. Default 0.
+    pin_memory : bool
+        If True, use pinned memory for faster GPU transfers. This increases memory
+        usage significantly (roughly 2x the prefetch buffer size) but can speed up
+        data transfer to GPU. Default is True.
+    prefetch_factor : int
+        Number of batches each worker prefetches. Higher values use more memory
+        but can improve throughput. Default is 1.
 
     Returns
     -------
@@ -215,22 +225,26 @@ def build_train_and_test_loaders(
         dataset, train_fraction
     )
 
+    # prefetch_factor is only valid when num_workers > 0
+    loader_kwargs = dict(
+        batch_size=batch_size,
+        pin_memory=pin_memory,
+        num_workers=num_workers,
+        worker_init_fn=fix_random_seeds,
+    )
+    if num_workers > 0:
+        loader_kwargs["prefetch_factor"] = prefetch_factor
+
     # Build DataLoaders
     train_loader = DataLoader(
         train_dataset,
-        batch_size=batch_size,
         shuffle=True,
-        pin_memory=True,
-        num_workers=num_workers,
-        worker_init_fn=fix_random_seeds,
+        **loader_kwargs,
     )
     test_loader = DataLoader(
         test_dataset,
-        batch_size=batch_size,
         shuffle=False,
-        pin_memory=True,
-        num_workers=num_workers,
-        worker_init_fn=fix_random_seeds,
+        **loader_kwargs,
     )
 
     return train_loader, test_loader

dingo/gw/training/train_pipeline.py

@@ -32,7 +32,10 @@
 
 
 def copy_files_to_local(
-    file_path: str, local_dir: Optional[str], leave_keys_on_disk: bool, is_condor: bool = False,
+    file_path: str,
+    local_dir: Optional[str],
+    leave_keys_on_disk: bool,
+    is_condor: bool = False,
 ) -> str:
     """
     Copy files to local node if local_dir is provided to minimize network traffic during training.
@@ -232,7 +235,7 @@ def initialize_stage(
     pm: BasePosteriorModel,
     wfd: WaveformDataset,
     stage: dict,
-    num_workers: int,
+    local_settings: dict,
     resume: bool = False,
 ):
     """
@@ -249,7 +252,8 @@ def initialize_stage(
     wfd : WaveformDataset
     stage : dict
         Settings specific to current stage of training
-    num_workers : int
+    local_settings : dict
+        Local settings for training (num_workers, pin_memory, prefetch_factor, etc.)
     resume : bool
         Whether training is resuming mid-stage. This controls whether the optimizer and
         scheduler should be re-initialized based on contents of stage dict.
@@ -269,7 +273,9 @@ def initialize_stage(
         wfd,
         train_settings["data"]["train_fraction"],
         stage["batch_size"],
-        num_workers,
+        num_workers=local_settings.get("num_workers", 0),
+        pin_memory=local_settings.get("pin_memory", True),
+        prefetch_factor=local_settings.get("prefetch_factor", 1),
     )
 
     if not resume:
@@ -343,13 +349,13 @@ def train_stages(
             print(f"\nBeginning training stage {n}. Settings:")
             print(yaml.dump(stage, default_flow_style=False, sort_keys=False))
             train_loader, test_loader = initialize_stage(
-                pm, wfd, stage, local_settings["num_workers"], resume=False
+                pm, wfd, stage, local_settings, resume=False
             )
         else:
             print(f"\nResuming training in stage {n}. Settings:")
             print(yaml.dump(stage, default_flow_style=False, sort_keys=False))
             train_loader, test_loader = initialize_stage(
-                pm, wfd, stage, local_settings["num_workers"], resume=True
+                pm, wfd, stage, local_settings, resume=True
             )
         early_stopping = None
         if stage.get("early_stopping"):

docs/source/training.md

@@ -102,6 +102,8 @@ local:
   checkpoint_epochs: 10
   leave_waveforms_on_disk: True
   local_cache_path: tmp
+  # pin_memory: True
+  # prefetch_factor: 1
 #   condor:
 #     bid: 100
 #     num_cpus: 16
@@ -172,6 +174,12 @@ leave_waveforms_on_disk
 local_cache_path
 : When training on a cluster and loading waveforms during training (i.e., `leave_waveforms_on_disk=True`), the waveform dataset should be copied to the disk storage of the local node at the beginning of training. This prevents unexpected long data loading times during training due to network traffic. Usually, paths for local storage are `tmp` or `dev/shm`. When submitting the job with `condor`, `request_disk: 50GB` should be included in the `condor` settings with the requested disk space larger than the size of the waveform dataset used for training.
 
+pin_memory
+: If `True` (default), use pinned (page-locked) memory for faster CPU to GPU data transfers. This speeds up training by ~20% but increases memory usage. The memory overhead scales as `num_workers × prefetch_factor × batch_size × sample_size`, and is roughly doubled when using pinned memory (once for the worker prefetch queue, once for the pinned copy). Set to `False` if memory is constrained.
+
+prefetch_factor
+: Number of batches each DataLoader worker prefetches (default 1). Higher values can improve throughput by hiding data loading latency, but increase memory usage proportionally. In testing, `prefetch_factor=1` performs as well as `prefetch_factor=2` while using less memory.
+
 condor
 : Settings for [HTCondor](https://htcondor.readthedocs.io/en/latest/index.html). The condor script will (re)submit itself according to these options.
 

examples/fmpe_model/train_settings.yaml

@@ -118,3 +118,5 @@ local:
     max_epochs_per_run: 500
     max_time_per_run: 360000
   leave_waveforms_on_disk: True
+  # pin_memory: True  # Set False to reduce memory usage, but slows down transfer of batches from CPU to GPU
+  # prefetch_factor: 1  # Number of batches each worker prefetches; higher uses more memory

examples/gnpe_model/train_settings.yaml

@@ -111,6 +111,8 @@ local:
     max_epochs_per_run: 500
   checkpoint_epochs: 10
   leave_waveforms_on_disk: True
+  # pin_memory: True  # Set False to reduce memory usage, but slows down transfer of batches from CPU to GPU
+  # prefetch_factor: 1  # Number of batches each worker prefetches; higher uses more memory
   # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster
   # Local settings related to condor, remove if not used on cluster
   condor:

examples/gnpe_model/train_settings_init.yaml

@@ -90,6 +90,8 @@ local:
     max_epochs_per_run: 500
   checkpoint_epochs: 10
   leave_waveforms_on_disk: True
+  # pin_memory: True  # Set False to reduce memory usage, but slows down transfer of batches from CPU to GPU
+  # prefetch_factor: 1  # Number of batches each worker prefetches; higher uses more memory
   # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster
   condor:
     num_cpus: 16

examples/npe_model/train_settings.yaml

@@ -110,6 +110,8 @@ local:
     max_epochs_per_run: 500
   checkpoint_epochs: 50
   leave_waveforms_on_disk: True
+  # pin_memory: True  # Set False to reduce memory usage, but slows down transfer of batches from CPU to GPU
+  # prefetch_factor: 1  # Number of batches each worker prefetches; higher uses more memory
   # local_cache_path: tmp # uncomment to avoid slow data loading during training due to network traffic on cluster
   # Local settings related to condor, remove if not used on cluster
   condor:

examples/toy_npe_model/train_settings.yaml

@@ -76,4 +76,6 @@ local:
   runtime_limits:
     max_time_per_run: 3600000
     max_epochs_per_run: 30
-  checkpoint_epochs: 15
+  checkpoint_epochs: 15
+  # pin_memory: True  # Set False to reduce memory usage, but slows down transfer of batches from CPU to GPU
+  # prefetch_factor: 1  # Number of batches each worker prefetches; higher uses more memory