From 67e223b5028fac08c7266ef6e1e30f01330fbeb3 Mon Sep 17 00:00:00 2001 From: Booyaka101 Date: Fri, 1 May 2026 23:07:43 +0800 Subject: [PATCH] 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. --- toolkit/data_loader.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index 1d5536e..8f6fb81 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -694,6 +694,13 @@ 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(): + dataloader_kwargs['pin_memory'] = True + if has_buckets: # make sure they all have buckets for dataset in datasets: