Visitar URL original
Merge pull request #3550 from lu123bing/feat/dataloader-throughput-se… · DeepLabCut/DeepLabCut@2920414 · GitHub
Skip to content

Commit 2920414

Browse files
deruyter92deruyter92
authored
Merge pull request #3550 from lu123bing/feat/dataloader-throughput-settings-3547
Expose DataLoader persistent_workers and prefetch_factor in train_settings
2 parents 265d063 + 40905e1 commit 2920414

7 files changed

Lines changed: 211 additions & 2 deletions

File tree

‎deeplabcut/core/config/validation.py‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from collections.abc import Sequence
1+
from collections.abc import Mapping, Sequence
22
from typing import Annotated, Any, TypeAlias, TypeVar
33

44
import numpy as np
@@ -54,6 +54,40 @@ def validate_crop_bounds(
5454
less_than(y1, y2, name="y1", threshold_name="y2")
5555

5656

57+
def reject_reserved_keys(
58+
kwargs: Mapping[str, Any] | None,
59+
reserved: Mapping[str, str | None],
60+
name: str,
61+
) -> None:
62+
"""Rejects reserved keys in a pass-through kwargs mapping.
63+
64+
Example:
65+
```python
66+
class ExampleDataLoaderConfig(DLCBaseConfig):
67+
batch_size: int = 8
68+
dataloader_kwargs: dict[str, Any] | None = None
69+
70+
@field_validator("dataloader_kwargs")
71+
@classmethod
72+
def validate_dataloader_kwargs(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
73+
reject_reserved_keys(v, {"batch_size": "batch_size"}, name="dataloader_kwargs")
74+
return v
75+
```
76+
77+
Args:
78+
kwargs: The pass-through kwargs to check (None is treated as empty).
79+
reserved: The reserved keys that are forbidden in the pass-through kwargs.
80+
name: Name of the pass-through field, used in the error message.
81+
82+
Raises:
83+
ValueError: If ``kwargs`` contains any reserved key.
84+
"""
85+
clashes = sorted(set(kwargs or {}) & set(reserved))
86+
if clashes:
87+
hints = [f"'{key}' (set via `{reserved[key]}`)" if reserved[key] else f"'{key}'" for key in clashes]
88+
raise ValueError(f"{name} cannot include reserved options: {', '.join(hints)}")
89+
90+
5791
def _coerce_ndarray(v):
5892
if isinstance(v, np.ndarray):
5993
return v

‎deeplabcut/pose_estimation_pytorch/apis/training.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,13 +153,15 @@ def train(
153153
batch_size = run_config["train_settings"]["batch_size"]
154154
num_workers = run_config["train_settings"]["dataloader_workers"]
155155
pin_memory = run_config["train_settings"]["dataloader_pin_memory"]
156+
dataloader_kwargs = run_config["train_settings"].get("dataloader_kwargs") or {}
156157
train_dataloader = DataLoader(
157158
train_dataset,
158159
batch_size=batch_size,
159160
shuffle=True,
160161
collate_fn=collate_fn,
161162
num_workers=num_workers,
162163
pin_memory=pin_memory,
164+
**dataloader_kwargs,
163165
)
164166
valid_dataloader = DataLoader(valid_dataset, batch_size=1, shuffle=False)
165167

‎deeplabcut/pose_estimation_pytorch/config/training.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,24 @@
1010
#
1111
"""Training configuration classes for DeepLabCut pose estimation models."""
1212

13+
from typing import Any
14+
15+
from pydantic import field_validator
16+
1317
from deeplabcut.core.config import DLCBaseConfig
18+
from deeplabcut.core.config.validation import reject_reserved_keys
1419
from deeplabcut.core.weight_init import WeightInitialization
1520

21+
# DataLoader arguments set by DeepLabCut, mapped to the option that controls them
22+
_RESERVED_DATALOADER_KWARGS = {
23+
"dataset": None,
24+
"batch_size": "batch_size",
25+
"shuffle": None,
26+
"collate_fn": "data.train.collate",
27+
"num_workers": "dataloader_workers",
28+
"pin_memory": "dataloader_pin_memory",
29+
}
30+
1631

1732
class TrainSettingsConfig(DLCBaseConfig):
1833
"""Training settings configuration.
@@ -21,6 +36,9 @@ class TrainSettingsConfig(DLCBaseConfig):
2136
batch_size: Training batch size
2237
dataloader_workers: Number of data loader workers
2338
dataloader_pin_memory: Whether to pin memory in data loader
39+
dataloader_kwargs: Additional keyword arguments for the training
40+
``torch.utils.data.DataLoader`` (e.g. ``persistent_workers``). Arguments
41+
that DeepLabCut sets itself are reserved and cannot be given here.
2442
display_iters: Display interval for training progress
2543
epochs: Number of training epochs
2644
seed: Random seed for reproducibility
@@ -30,9 +48,16 @@ class TrainSettingsConfig(DLCBaseConfig):
3048
batch_size: int = 8
3149
dataloader_workers: int = 0
3250
dataloader_pin_memory: bool = False
51+
dataloader_kwargs: dict[str, Any] | None = None
3352
display_iters: int = 500
3453
epochs: int = 200
3554
seed: int = 42
3655
# @TODO @deruyter92 2026-02-13: V0 pipeline uses None for default initialization
3756
# (with ImageNet weights). We should update this to explicit WeightInitialization.
3857
weight_init: WeightInitialization | None = None
58+
59+
@field_validator("dataloader_kwargs")
60+
@classmethod
61+
def validate_dataloader_kwargs(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
62+
reject_reserved_keys(v, _RESERVED_DATALOADER_KWARGS, name="dataloader_kwargs")
63+
return v

‎docs/pytorch/pytorch_config.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -439,11 +439,19 @@ and [memory pinning](https://pytorch.org/docs/stable/data.html#memory-pinning).
439439
will use multi-process data loading. You should always keep
440440
`dataloader_pin_memory: true` when training on an NVIDIA GPU.
441441

442+
Any other [`DataLoader`](https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader)
443+
argument can be passed to the training data loader with `dataloader_kwargs`. For example,
444+
`persistent_workers: true` keeps the workers alive between epochs, and `prefetch_factor`
445+
sets how many batches each worker loads in advance (both require `dataloader_workers > 0`).
446+
Arguments that DeepLabCut sets itself (`dataset`, `batch_size`, `shuffle`, `collate_fn`,
447+
`num_workers` and `pin_memory`) are reserved and cannot be set in `dataloader_kwargs`.
448+
442449
```yaml
443450
train_settings:
444451
batch_size: 1 # the batch size used for training
445452
dataloader_workers: 0 # the number of workers for the PyTorch Dataloader
446453
dataloader_pin_memory: true # pin DataLoader memory
454+
dataloader_kwargs: null # extra DataLoader arguments, e.g. {persistent_workers: true}
447455
display_iters: 500 # the number of iterations (steps) between each log print
448456
epochs: 200 # the maximum number of epochs for which to train the model
449457
seed: 42 # the random seed to set for reproducibility

‎tests/core/config/test_validation.py‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,11 @@
77
#
88
"""Tests for annotated types in deeplabcut.core.config.validation."""
99

10+
from typing import Any
11+
1012
import numpy as np
1113
import pytest
12-
from pydantic import ValidationError
14+
from pydantic import ValidationError, field_validator
1315

1416
from deeplabcut.core.config import DLCBaseConfig
1517
from deeplabcut.core.config.validation import (
@@ -20,6 +22,7 @@
2022
UniqueStrList,
2123
greater_than,
2224
less_than,
25+
reject_reserved_keys,
2326
unique_values,
2427
validate_crop_bounds,
2528
)
@@ -40,6 +43,19 @@ class ListConfig(DLCBaseConfig):
4043
indices: NDArrayInt = None
4144

4245

46+
class KwargsConfig(DLCBaseConfig):
47+
"""Config exercising a pass-through kwargs field with reserved keys."""
48+
49+
name: str = "run"
50+
call_kwargs: dict[str, Any] | None = None
51+
52+
@field_validator("call_kwargs")
53+
@classmethod
54+
def validate_call_kwargs(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
55+
reject_reserved_keys(v, {"name": "name", "dataset": None}, name="call_kwargs")
56+
return v
57+
58+
4359
# ------------------------------------------------------------------
4460
# DefaultIfNone
4561
# ------------------------------------------------------------------
@@ -215,3 +231,42 @@ def test_existing_array_is_passed_through_unchanged(self):
215231
indices = ListConfig(indices=array).indices
216232
assert indices is array
217233
assert indices.dtype == np.float64
234+
235+
236+
# ------------------------------------------------------------------
237+
# reject_reserved_keys
238+
# ------------------------------------------------------------------
239+
240+
241+
class TestValidateNoReservedKeys:
242+
@pytest.mark.parametrize("value", [None, {}, {"mode": "offline", "tags": ["a"]}], ids=["none", "empty", "free"])
243+
def test_accepts_unreserved_kwargs(self, value):
244+
assert KwargsConfig(call_kwargs=value).call_kwargs == value
245+
246+
def test_rejects_reserved_key_with_field_hint(self):
247+
with pytest.raises(
248+
ValidationError, match=r"call_kwargs cannot include reserved options: 'name' \(set via `name`\)"
249+
):
250+
KwargsConfig(call_kwargs={"name": "other"})
251+
252+
def test_rejects_reserved_key_even_when_value_matches_field(self):
253+
with pytest.raises(ValidationError, match="reserved options"):
254+
KwargsConfig(name="run", call_kwargs={"name": "run"})
255+
256+
def test_rejects_reserved_key_without_config_field(self):
257+
with pytest.raises(ValidationError, match="reserved options: 'dataset'") as excinfo:
258+
KwargsConfig(call_kwargs={"dataset": "x"})
259+
assert "set via" not in str(excinfo.value)
260+
261+
def test_rejects_reserved_key_on_assignment(self):
262+
cfg = KwargsConfig()
263+
with pytest.raises(ValidationError, match="reserved options"):
264+
cfg.call_kwargs = {"dataset": "x"}
265+
266+
@pytest.mark.parametrize("kwargs", [None, {}, {"mode": "offline"}], ids=["none", "empty", "free"])
267+
def test_function_accepts_unreserved_kwargs(self, kwargs):
268+
reject_reserved_keys(kwargs, {"dataset": None}, name="kw")
269+
270+
def test_function_lists_all_clashes_sorted(self):
271+
with pytest.raises(ValueError, match=r"kw cannot include reserved options: 'a' \(set via `x`\), 'b'$"):
272+
reject_reserved_keys({"b": 1, "a": 2, "c": 3}, {"a": "x", "b": None}, name="kw")

‎tests/pose_estimation_pytorch/apis/test_apis_training.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,3 +55,34 @@ def test_train_uses_resume_training_from_config(
5555
train(loader=loader, run_config=run_config, task=Task.BOTTOM_UP, device="cpu", snapshot_path=None)
5656

5757
assert mock_build_runner.call_args.kwargs["snapshot_path"] == "/train/snapshot-010.pt"
58+
59+
60+
@patch("deeplabcut.pose_estimation_pytorch.apis.training.DataLoader")
61+
@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_transforms", return_value=Mock())
62+
@patch("deeplabcut.pose_estimation_pytorch.apis.training.PoseModel.build", return_value=Mock())
63+
@patch("deeplabcut.pose_estimation_pytorch.apis.training.build_training_runner", return_value=Mock())
64+
def test_train_passes_dataloader_kwargs_to_train_dataloader(
65+
mock_build_runner: Mock,
66+
mock_build_model: Mock,
67+
mock_build_transforms: Mock,
68+
mock_dataloader: Mock,
69+
tmp_path: Path,
70+
) -> None:
71+
run_config = _minimal_run_config(tmp_path)
72+
run_config.train_settings.dataloader_workers = 2
73+
run_config.train_settings.dataloader_kwargs = {"persistent_workers": True, "prefetch_factor": 4}
74+
75+
loader = Mock()
76+
loader.model_folder = tmp_path
77+
loader.model_cfg = run_config
78+
train_dataset = Mock(__len__=Mock(return_value=1))
79+
valid_dataset = Mock(__len__=Mock(return_value=1))
80+
loader.create_dataset = Mock(side_effect=[train_dataset, valid_dataset])
81+
82+
train(loader=loader, run_config=run_config, task=Task.BOTTOM_UP, device="cpu", snapshot_path=None)
83+
84+
train_call, valid_call = mock_dataloader.call_args_list
85+
assert train_call.kwargs["num_workers"] == 2
86+
assert train_call.kwargs["persistent_workers"] is True
87+
assert train_call.kwargs["prefetch_factor"] == 4
88+
assert "persistent_workers" not in valid_call.kwargs

‎tests/pose_estimation_pytorch/config/test_train_settings_config.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
#
2+
# DeepLabCut Toolbox (deeplabcut.org)
3+
# © A. & M.W. Mathis Labs
4+
# https://github.com/DeepLabCut/DeepLabCut
5+
#
6+
# Please see AUTHORS for contributors.
7+
# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS
8+
#
9+
# Licensed under GNU Lesser General Public License v3.0
10+
#
11+
"""Tests for TrainSettingsConfig validation."""
12+
13+
from __future__ import annotations
14+
15+
import pytest
16+
from pydantic import ValidationError
17+
18+
from deeplabcut.pose_estimation_pytorch.config.training import TrainSettingsConfig
19+
20+
21+
def test_dataloader_kwargs_defaults_to_none() -> None:
22+
assert TrainSettingsConfig().dataloader_kwargs is None
23+
24+
25+
def test_old_configs_without_dataloader_kwargs_still_validate() -> None:
26+
cfg = TrainSettingsConfig.from_dict({"batch_size": 4, "dataloader_workers": 2})
27+
assert cfg.dataloader_kwargs is None
28+
29+
30+
def test_dataloader_kwargs_accepts_extra_options() -> None:
31+
kwargs = {"persistent_workers": True, "prefetch_factor": 4}
32+
cfg = TrainSettingsConfig.from_dict({"dataloader_workers": 2, "dataloader_kwargs": kwargs})
33+
assert cfg.dataloader_kwargs == kwargs
34+
assert cfg.to_dict(normalize=True)["dataloader_kwargs"] == kwargs
35+
36+
37+
@pytest.mark.parametrize(
38+
("key", "hint"),
39+
[
40+
("num_workers", "dataloader_workers"),
41+
("pin_memory", "dataloader_pin_memory"),
42+
("batch_size", "batch_size"),
43+
("collate_fn", "data.train.collate"),
44+
],
45+
)
46+
def test_dataloader_kwargs_rejects_options_with_a_config_field(key: str, hint: str) -> None:
47+
with pytest.raises(ValidationError, match=f"'{key}' \\(set via `{hint}`\\)"):
48+
TrainSettingsConfig.from_dict({"dataloader_kwargs": {key: 1}})
49+
50+
51+
@pytest.mark.parametrize("key", ["dataset", "shuffle"])
52+
def test_dataloader_kwargs_rejects_options_managed_by_deeplabcut(key: str) -> None:
53+
with pytest.raises(ValidationError, match=f"dataloader_kwargs cannot include reserved options: '{key}'"):
54+
TrainSettingsConfig.from_dict({"dataloader_kwargs": {key: False}})

0 commit comments

Comments
 (0)