diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 5ed490e..4c95909 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -977,6 +977,12 @@ class DatasetConfig: self.num_workers: int = kwargs.get('num_workers', 2) self.prefetch_factor: int = kwargs.get('prefetch_factor', 2) + # Pin DataLoader output tensors in page-locked RAM for faster CPU->GPU + # transfer. Off by default because page-locked RAM cannot be relocated + # by NVIDIA's Windows driver shared-memory VRAM-overflow fallback, + # which can cause severe PCIe thrashing for users at the VRAM ceiling. + # Opt in if you have stable VRAM headroom and want the transfer speedup. + self.pin_memory: bool = kwargs.get('pin_memory', False) self.extra_values: List[float] = kwargs.get('extra_values', []) self.square_crop: bool = kwargs.get('square_crop', False) # apply same augmentations to control images. Usually want this true unless special case diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 8f6fb81..6b39e26 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -694,11 +694,14 @@ def get_dataloader_from_datasets( dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor - # Enable pinned memory for faster CPU->GPU transfer when CUDA is available. - # pin_memory is independent of num_workers and works even with num_workers=0 - # (the forced setting on native Windows / macOS due to dataset pickling - # constraints). No-op on non-CUDA backends. - if torch.cuda.is_available(): + # Enable pinned memory only when explicitly opted in via dataset config and + # CUDA is available. pin_memory speeds up CPU->GPU transfer (helpful even at + # num_workers=0), but page-locked RAM cannot be relocated by NVIDIA's + # Windows driver shared-memory VRAM-overflow fallback, which can cause + # severe PCIe thrashing for users at the VRAM ceiling. Off by default; opt + # in via 'pin_memory: true' on the dataset config when VRAM headroom is + # stable. + if torch.cuda.is_available() and dataset_config_list[0].pin_memory: dataloader_kwargs['pin_memory'] = True if has_buckets: