Expose persistent_workers and prefetch_factor for the PyTorch training DataLoader #3547

Is your feature request related to a problem? Please describe.

DeepLabCut 3.0.2 exposes DataLoader worker count and pinned-memory settings, but not PyTorch's persistent_workers or prefetch_factor.

In my multi-animal project, training shows bursty GPU activity and appears limited by the input pipeline. My immediate motivation is an HPC policy requiring at least 75% average GPU utilization, but this is a general throughput and configurability issue for projects with CPU-heavy decoding, augmentation, short epochs, or large GPUs.

The project uses 1,225 training frames, PNG inputs staged on node-local NVMe, 736 × 736 keypoint-aware crops, rotation, scaling, and histogram equalization. Pinned memory was enabled throughout.

A representative DEKR-W48 run used batch size 32 and four DataLoader workers:

Allocated CPU cores: 32
Wall time:           11:35:56
CPU time:            16:53:44
System/I/O:             27.27%
Disk read:            466.66 GiB
Average GPU use:          68%

Tests of the currently exposed controls did not provide a reliable solution:

Batch 16, 6 workers: 56% average GPU utilization
Batch 32, 6 workers: 55% average GPU utilization
Batch 48 and 64:      CUDA out of memory
28 workers:           substantially slower with lower GPU utilization

With two workers and 32 allocated CPU cores, process monitoring showed:

Main process: 98.5% CPU
Worker 1:     27.0% CPU
Worker 2:     25.9% CPU

Additional allocated CPUs therefore remained largely unused, while increasing the worker count reduced throughput.

DeepLabCut currently constructs the training DataLoader without specifying persistent_workers or prefetch_factor:

train_dataloader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=collate_fn,
    num_workers=num_workers,
    pin_memory=pin_memory,
)

This leaves the PyTorch defaults in effect:

persistent_workers=False
prefetch_factor=2

For relatively small datasets trained over hundreds of epochs, repeated worker creation may add avoidable overhead. Configurable prefetching would also allow queue depth to be tuned independently of worker count.

This is more specific than issues #2343 and #2639. Increasing batch size, the main recommendation in earlier discussions, did not improve average GPU utilization here and was ultimately limited by GPU memory.

Describe the solution you'd like

Please expose persistent_workers and prefetch_factor through TrainSettingsConfig, pytorch_config.yaml, and pytorch_cfg_updates.

For example:

train_settings:
  dataloader_workers: 4
  dataloader_pin_memory: true
  dataloader_persistent_workers: true
  dataloader_prefetch_factor: 2

The exact names could follow DeepLabCut's existing conventions. These settings should be validated and persisted in the effective configuration for reproducibility.

It would also be helpful to document whether pinned-memory batches are transferred to CUDA using non_blocking=True.

Describe alternatives you've considered

I tested:

• increasing batch size up to the GPU-memory limit;
• varying DataLoader worker count;
• allocating additional CPU cores;
• limiting numerical-library threads;
• staging training inputs on node-local NVMe storage.

Increasing batch size did not improve average utilization, and higher worker counts reduced throughput. Removing CPU-side augmentations may improve throughput, but could compromise model quality.

Patching or monkey-patching torch.utils.data.DataLoader would be fragile across DeepLabCut updates and would bypass the validated configuration system.

Additional context

Environment:

• DeepLabCut 3.0.2
• PyTorch backend
• Multi-animal project
• Linux and Slurm
• NVIDIA H200 with 141 GB VRAM
• Conda environment
• Training inputs staged on node-local NVMe

I can provide complete training logs, effective configurations, resource summaries, and controlled benchmark results if helpful.

Here nvtop outputs from some of the jobs:
Batch size 16, 8 cores

Batch size 32, 16 cores, 14 workers

Batch size 32, 32 cores, 14 workers

Batch size 32. 32 cores, 4 workers

Batch size 32, 32 cores, 2 workers

Disclosure: Microsoft 365 Copilot assisted with structuring and editing this issue. The author verified the reported observations and remains responsible for the content.