deruyter92 · Oct 13, 2026 · Oct 13, 2026 · Oct 16, 2026 · Oct 16, 2026 · Oct 16, 2026
diff --git a/deeplabcut/core/config/validation.py b/deeplabcut/core/config/validation.py
@@ -1,4 +1,4 @@
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from typing import Annotated, Any, TypeAlias, TypeVar
 
 import numpy as np
@@ -54,6 +54,40 @@ def validate_crop_bounds(
     less_than(y1, y2, name="y1", threshold_name="y2")
 
 
+def reject_reserved_keys(
+    kwargs: Mapping[str, Any] | None,
+    reserved: Mapping[str, str | None],
+    name: str,
+) -> None:
+    """Rejects reserved keys in a pass-through kwargs mapping.
+
+    Example:
+        ```python
+        class ExampleDataLoaderConfig(DLCBaseConfig):
+            batch_size: int = 8
+            dataloader_kwargs: dict[str, Any] | 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, {"batch_size": "batch_size"}, name="dataloader_kwargs")
+                return v
+        ```
+
+    Args:
+        kwargs: The pass-through kwargs to check (None is treated as empty).
+        reserved: The reserved keys that are forbidden in the pass-through kwargs.
+        name: Name of the pass-through field, used in the error message.
+
+    Raises:
+        ValueError: If ``kwargs`` contains any reserved key.
+    """
+    clashes = sorted(set(kwargs or {}) & set(reserved))
+    if clashes:
+        hints = [f"'{key}' (set via `{reserved[key]}`)" if reserved[key] else f"'{key}'" for key in clashes]
+        raise ValueError(f"{name} cannot include reserved options: {', '.join(hints)}")
+
+
 def _coerce_ndarray(v):
     if isinstance(v, np.ndarray):
         return v

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/core/config/test_validation.py b/tests/core/config/test_validation.py
@@ -7,9 +7,11 @@
 #
 """Tests for annotated types in deeplabcut.core.config.validation."""
 
+from typing import Any
+
 import numpy as np
 import pytest
-from pydantic import ValidationError
+from pydantic import ValidationError, field_validator
 
 from deeplabcut.core.config import DLCBaseConfig
 from deeplabcut.core.config.validation import (
@@ -20,6 +22,7 @@
     UniqueStrList,
     greater_than,
     less_than,
+    reject_reserved_keys,
     unique_values,
     validate_crop_bounds,
 )
@@ -40,6 +43,19 @@ class ListConfig(DLCBaseConfig):
     indices: NDArrayInt = None
 
 
+class KwargsConfig(DLCBaseConfig):
+    """Config exercising a pass-through kwargs field with reserved keys."""
+
+    name: str = "run"
+    call_kwargs: dict[str, Any] | None = None
+
+    @field_validator("call_kwargs")
+    @classmethod
+    def validate_call_kwargs(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
+        reject_reserved_keys(v, {"name": "name", "dataset": None}, name="call_kwargs")
+        return v
+
+
 # ------------------------------------------------------------------
 # DefaultIfNone
 # ------------------------------------------------------------------
@@ -215,3 +231,42 @@ def test_existing_array_is_passed_through_unchanged(self):
         indices = ListConfig(indices=array).indices
         assert indices is array
         assert indices.dtype == np.float64
+
+
+# ------------------------------------------------------------------
+# reject_reserved_keys
+# ------------------------------------------------------------------
+
+
+class TestValidateNoReservedKeys:
+    @pytest.mark.parametrize("value", [None, {}, {"mode": "offline", "tags": ["a"]}], ids=["none", "empty", "free"])
+    def test_accepts_unreserved_kwargs(self, value):
+        assert KwargsConfig(call_kwargs=value).call_kwargs == value
+
+    def test_rejects_reserved_key_with_field_hint(self):
+        with pytest.raises(
+            ValidationError, match=r"call_kwargs cannot include reserved options: 'name' \(set via `name`\)"
+        ):
+            KwargsConfig(call_kwargs={"name": "other"})
+
+    def test_rejects_reserved_key_even_when_value_matches_field(self):
+        with pytest.raises(ValidationError, match="reserved options"):
+            KwargsConfig(name="run", call_kwargs={"name": "run"})
+
+    def test_rejects_reserved_key_without_config_field(self):
+        with pytest.raises(ValidationError, match="reserved options: 'dataset'") as excinfo:
+            KwargsConfig(call_kwargs={"dataset": "x"})
+        assert "set via" not in str(excinfo.value)
+
+    def test_rejects_reserved_key_on_assignment(self):
+        cfg = KwargsConfig()
+        with pytest.raises(ValidationError, match="reserved options"):
+            cfg.call_kwargs = {"dataset": "x"}
+
+    @pytest.mark.parametrize("kwargs", [None, {}, {"mode": "offline"}], ids=["none", "empty", "free"])
+    def test_function_accepts_unreserved_kwargs(self, kwargs):
+        reject_reserved_keys(kwargs, {"dataset": None}, name="kw")
+
+    def test_function_lists_all_clashes_sorted(self):
+        with pytest.raises(ValueError, match=r"kw cannot include reserved options: 'a' \(set via `x`\), 'b'$"):
+            reject_reserved_keys({"b": 1, "a": 2, "c": 3}, {"a": "x", "b": None}, name="kw")
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}})