perf: enable pin_memory on CUDA DataLoaders for faster CPU->GPU transfer
torch DataLoader's pin_memory was never set, defaulting to False. Setting
pin_memory=True allocates batch tensors in page-locked memory, which lets
PyTorch issue async cuda host->device copies via DMA - typically 10-30%
faster transfer for image/video training tensors and can reduce per-step
overhead noticeably on Windows where num_workers is forced to 0 (see
6d31c6db / #770) and the main process pays the full transfer cost.
pin_memory is orthogonal to num_workers and is safe with num_workers=0,
so this benefits Windows/macOS users (who currently can't use multi-worker
loading due to pickling constraints) without changing any other behavior.
Gated on torch.cuda.is_available() so it's a no-op on CPU/MPS/ROCm-without
-CUDA-API setups.
Refs #758.
This commit is contained in:
1 parent
0d91fcee9e
commit
67e223b502
1 file changed
+7
@@ -694,6 +694,13 @@ def get_dataloader_from_datasets(
|
|||||||
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
|
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
|
||||||
dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor
|
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():
|
||||||
|
dataloader_kwargs['pin_memory'] = True
|
||||||
|
|
||||||
if has_buckets:
|
if has_buckets:
|
||||||
# make sure they all have buckets
|
# make sure they all have buckets
|
||||||
for dataset in datasets:
|
for dataset in datasets:
|
||||||
|
|||||||
Reference in new issue
Block a user