add TrainSettingsConfig.dataloader_kwargs

deruyter92 · Oct 13, 2026 · Oct 13, 2026 · Oct 16, 2026 · Oct 16, 2026 · Oct 16, 2026
commit 53c5b94e25d5bc8f02f98b5016b76be1fc5312db
diff --git a/deeplabcut/pose_estimation_pytorch/apis/training.py b/deeplabcut/pose_estimation_pytorch/apis/training.py
@@ -153,13 +153,15 @@ def train(
     batch_size = run_config["train_settings"]["batch_size"]
     num_workers = run_config["train_settings"]["dataloader_workers"]
     pin_memory = run_config["train_settings"]["dataloader_pin_memory"]
+    dataloader_kwargs = run_config["train_settings"].get("dataloader_kwargs") or {}
     train_dataloader = DataLoader(
         train_dataset,
         batch_size=batch_size,
         shuffle=True,
         collate_fn=collate_fn,
         num_workers=num_workers,
         pin_memory=pin_memory,
+        **dataloader_kwargs,
     )
     valid_dataloader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
 

diff --git a/deeplabcut/pose_estimation_pytorch/config/training.py b/deeplabcut/pose_estimation_pytorch/config/training.py
@@ -10,9 +10,24 @@
 #
 """Training configuration classes for DeepLabCut pose estimation models."""
 
+from typing import Any
+
+from pydantic import field_validator
+
 from deeplabcut.core.config import DLCBaseConfig
+from deeplabcut.core.config.validation import reject_reserved_keys
 from deeplabcut.core.weight_init import WeightInitialization
 
+# DataLoader arguments set by DeepLabCut, mapped to the option that controls them
+_RESERVED_DATALOADER_KWARGS = {
+    "dataset": None,
+    "batch_size": "batch_size",
+    "shuffle": None,
+    "collate_fn": "data.collate",
+    "num_workers": "dataloader_workers",
+    "pin_memory": "dataloader_pin_memory",
+}
+
 
 class TrainSettingsConfig(DLCBaseConfig):
     """Training settings configuration.
@@ -21,6 +36,9 @@ class TrainSettingsConfig(DLCBaseConfig):
         batch_size: Training batch size
         dataloader_workers: Number of data loader workers
         dataloader_pin_memory: Whether to pin memory in data loader
+        dataloader_kwargs: Additional keyword arguments for the training
+            ``torch.utils.data.DataLoader`` (e.g. ``persistent_workers``). Arguments
+            that DeepLabCut sets itself are reserved and cannot be given here.
         display_iters: Display interval for training progress
         epochs: Number of training epochs
         seed: Random seed for reproducibility
@@ -30,9 +48,16 @@ class TrainSettingsConfig(DLCBaseConfig):
     batch_size: int = 8
     dataloader_workers: int = 0
     dataloader_pin_memory: bool = False
+    dataloader_kwargs: dict[str, Any] | None = None
     display_iters: int = 500
     epochs: int = 200
     seed: int = 42
     # @TODO @deruyter92 2026-02-13: V0 pipeline uses None for default initialization
     # (with ImageNet weights). We should update this to explicit WeightInitialization.
     weight_init: WeightInitialization | None = None
+
+    @field_validator("dataloader_kwargs")
+    @classmethod
+    def validate_dataloader_kwargs(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
+        reject_reserved_keys(v, _RESERVED_DATALOADER_KWARGS, name="dataloader_kwargs")
+        return v
diff --git a/docs/pytorch/pytorch_config.md b/docs/pytorch/pytorch_config.md
@@ -439,11 +439,19 @@ and [memory pinning](https://pytorch.org/docs/stable/data.html#memory-pinning).
 will use multi-process data loading. You should always keep
 `dataloader_pin_memory: true` when training on an NVIDIA GPU.
 
+Any other [`DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)
+argument can be passed to the training data loader with `dataloader_kwargs`. For example,
+`persistent_workers: true` keeps the workers alive between epochs, and `prefetch_factor`
+sets how many batches each worker loads in advance (both require `dataloader_workers > 0`).
+Arguments that DeepLabCut sets itself (`dataset`, `batch_size`, `shuffle`, `collate_fn`,
+`num_workers` and `pin_memory`) are reserved and cannot be set in `dataloader_kwargs`.
+
 ```yaml
 train_settings:
   batch_size: 1  # the batch size used for training
   dataloader_workers: 0  # the number of workers for the PyTorch Dataloader
   dataloader_pin_memory: true  # pin DataLoader memory
+  dataloader_kwargs: null  # extra DataLoader arguments, e.g. {persistent_workers: true}
   display_iters: 500  # the number of iterations (steps) between each log print
   epochs: 200  # the maximum number of epochs for which to train the model
   seed: 42  # the random seed to set for reproducibility

diff --git a/tests/pose_estimation_pytorch/apis/test_apis_training.py b/tests/pose_estimation_pytorch/apis/test_apis_training.py
@@ -55,3 +55,34 @@ def test_train_uses_resume_training_from_config(
     train(loader=loader, run_config=run_config, task=Task.BOTTOM_UP, device="cpu", snapshot_path=None)
 
     assert mock_build_runner.call_args.kwargs["snapshot_path"] == "/train/snapshot-010.pt"
+
+
+@patch("deeplabcut.pose_estimation_pytorch.apis.training.DataLoader")
+@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_transforms", return_value=Mock())
+@patch("deeplabcut.pose_estimation_pytorch.apis.training.PoseModel.build", return_value=Mock())
+@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_training_runner", return_value=Mock())
+def test_train_passes_dataloader_kwargs_to_train_dataloader(
+    mock_build_runner: Mock,
+    mock_build_model: Mock,
+    mock_build_transforms: Mock,
+    mock_dataloader: Mock,
+    tmp_path: Path,
+) -> None:
+    run_config = _minimal_run_config(tmp_path)
+    run_config.train_settings.dataloader_workers = 2
+    run_config.train_settings.dataloader_kwargs = {"persistent_workers": True, "prefetch_factor": 4}
+
+    loader = Mock()
+    loader.model_folder = tmp_path
+    loader.model_cfg = run_config
+    train_dataset = Mock(__len__=Mock(return_value=1))
+    valid_dataset = Mock(__len__=Mock(return_value=1))
+    loader.create_dataset = Mock(side_effect=[train_dataset, valid_dataset])
+
+    train(loader=loader, run_config=run_config, task=Task.BOTTOM_UP, device="cpu", snapshot_path=None)
+
+    train_call, valid_call = mock_dataloader.call_args_list
+    assert train_call.kwargs["num_workers"] == 2
+    assert train_call.kwargs["persistent_workers"] is True
+    assert train_call.kwargs["prefetch_factor"] == 4
+    assert "persistent_workers" not in valid_call.kwargs
diff --git a/tests/pose_estimation_pytorch/config/test_train_settings_config.py b/tests/pose_estimation_pytorch/config/test_train_settings_config.py
@@ -0,0 +1,54 @@
+#
+# DeepLabCut Toolbox (deeplabcut.org)
+# © A. & M.W. Mathis Labs
+# https://github.com/DeepLabCut/DeepLabCut
+#
+# Please see AUTHORS for contributors.
+# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS
+#
+# Licensed under GNU Lesser General Public License v3.0
+#
+"""Tests for TrainSettingsConfig validation."""
+
+from __future__ import annotations
+
+import pytest
+from pydantic import ValidationError
+
+from deeplabcut.pose_estimation_pytorch.config.training import TrainSettingsConfig
+
+
+def test_dataloader_kwargs_defaults_to_none() -> None:
+    assert TrainSettingsConfig().dataloader_kwargs is None
+
+
+def test_old_configs_without_dataloader_kwargs_still_validate() -> None:
+    cfg = TrainSettingsConfig.from_dict({"batch_size": 4, "dataloader_workers": 2})
+    assert cfg.dataloader_kwargs is None
+
+
+def test_dataloader_kwargs_accepts_extra_options() -> None:
+    kwargs = {"persistent_workers": True, "prefetch_factor": 4}
+    cfg = TrainSettingsConfig.from_dict({"dataloader_workers": 2, "dataloader_kwargs": kwargs})
+    assert cfg.dataloader_kwargs == kwargs
+    assert cfg.to_dict(normalize=True)["dataloader_kwargs"] == kwargs
+
+
+@pytest.mark.parametrize(
+    ("key", "hint"),
+    [
+        ("num_workers", "dataloader_workers"),
+        ("pin_memory", "dataloader_pin_memory"),
+        ("batch_size", "batch_size"),
+        ("collate_fn", "data.collate"),
+    ],
+)
+def test_dataloader_kwargs_rejects_options_with_a_config_field(key: str, hint: str) -> None:
+    with pytest.raises(ValidationError, match=f"'{key}' \\(set via `{hint}`\\)"):
+        TrainSettingsConfig.from_dict({"dataloader_kwargs": {key: 1}})
+
+
+@pytest.mark.parametrize("key", ["dataset", "shuffle"])
+def test_dataloader_kwargs_rejects_options_managed_by_deeplabcut(key: str) -> None:
+    with pytest.raises(ValidationError, match=f"dataloader_kwargs cannot include reserved options: '{key}'"):
+        TrainSettingsConfig.from_dict({"dataloader_kwargs": {key: False}})