diff --git a/dingo/core/utils/torchutils.py b/dingo/core/utils/torchutils.py index 3b37bca9f..619eaab2d 100644 --- a/dingo/core/utils/torchutils.py +++ b/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 diff --git a/dingo/gw/training/train_pipeline.py b/dingo/gw/training/train_pipeline.py index 57ca57061..5e9041a5d 100644 --- a/dingo/gw/training/train_pipeline.py +++ b/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"): diff --git a/docs/source/training.md b/docs/source/training.md index e48413fee..c22d80390 100644 --- a/docs/source/training.md +++ b/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. diff --git a/examples/fmpe_model/train_settings.yaml b/examples/fmpe_model/train_settings.yaml index 560c27127..72e64fdb6 100644 --- a/examples/fmpe_model/train_settings.yaml +++ b/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 diff --git a/examples/gnpe_model/train_settings.yaml b/examples/gnpe_model/train_settings.yaml index 9fd1d96eb..13a4cd246 100644 --- a/examples/gnpe_model/train_settings.yaml +++ b/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: diff --git a/examples/gnpe_model/train_settings_init.yaml b/examples/gnpe_model/train_settings_init.yaml index 676350392..4fb424389 100644 --- a/examples/gnpe_model/train_settings_init.yaml +++ b/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 diff --git a/examples/npe_model/train_settings.yaml b/examples/npe_model/train_settings.yaml index f6452b1e6..aabc9f632 100644 --- a/examples/npe_model/train_settings.yaml +++ b/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: diff --git a/examples/toy_npe_model/train_settings.yaml b/examples/toy_npe_model/train_settings.yaml index e331cdc78..5e933d39b 100644 --- a/examples/toy_npe_model/train_settings.yaml +++ b/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 \ No newline at end of file + 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 \ No newline at end of file