Skip to content

Commit d05dff7

Browse files
maximpavlivmaximpavliv
committed
Implement update_config_by_dotpath()
1 parent b2d2feb commit d05dff7

8 files changed

Lines changed: 152 additions & 45 deletions

File tree

‎deeplabcut/compat.py‎

Lines changed: 53 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,17 @@ def train_network(
7676
superanimal_name: str = "",
7777
superanimal_transfer_learning: bool = False,
7878
engine: Engine | None = None,
79-
**torch_kwargs,
79+
device: str | None = None,
80+
snapshot_path: str | Path | None = None,
81+
detector_path: str | Path | None = None,
82+
batch_size: int | None = None,
83+
detector_batch_size: int | None = None,
84+
detector_epochs: int | None = None,
85+
detector_save_epochs: int | None = None,
86+
pose_threshold: float | None = 0.1,
87+
pytorch_cfg_updates: dict | None = None,
8088
):
89+
8190
"""
8291
Trains the network with the labels in the training dataset.
8392
@@ -135,6 +144,7 @@ def train_network(
135144
Only for the PyTorch engine (equivalent to the `saveiters` parameter for the
136145
TensorFlow engine). The number of epochs between each snapshot save. If
137146
None, the value will be read from the `pytorch_config.yaml` file.
147+
138148
allow_growth: bool, optional, default=True.
139149
Only for the TensorFlow engine.
140150
For some smaller GPUs the memory issues happen. If ``True``, the memory
@@ -180,18 +190,38 @@ def train_network(
180190
overwrite this by passing the engine as an argument, but this should generally
181191
not be done.
182192
183-
torch_kwargs:
184-
You can add any keyword arguments for the deeplabcut.pose_estimation_pytorch
185-
train_network method here. These arguments are passed to the downstream method.
186-
Some of the parameters that can be passed are
187-
* ``device`` (the CUDA device to use for training)
188-
* ``batch_size`` (the batch size to use while training)
189-
* ``snapshot_path`` (the pose model snapshot to resume training from)
190-
* ``detector_path`` (the detector model snapshot to resume training from)
193+
device: str, optional, default = None.
194+
Only for the PyTorch engine. The device to run the training on (e.g. "cuda:0")
195+
196+
snapshot_path: str or Path, optional, default = None.
197+
Only for the PyTorch engine. The path to the pose model snapshot to resume training from.
198+
199+
detector_path: str or Path, optional, default = None.
200+
Only for the PyTorch engine. The path to the detector model snapshot to resume training from.
201+
202+
batch_size: int, optional, default = None.
203+
Only for the PyTorch engine. The batch size to use while training.
204+
205+
detector_batch_size: int, optional, default = None.
206+
Only for the PyTorch engine. The batch size to use while training the detector.
191207
192-
When training a top-down model, these parameters are also available for the
193-
detector, with the parameters ``detector_batch_size``, ``detector_epochs`` and
194-
``detector_save_epochs``.
208+
detector_epochs: int, optional, default = None.
209+
Only for the PyTorch engine. The number of epochs to train the detector for.
210+
211+
detector_save_epochs: int, optional, default = None.
212+
Only for the PyTorch engine. The number of epochs between each detector snapshot save.
213+
214+
pose_threshold: float, optional, default = 0.1.
215+
Only for the PyTorch engine. Used for memory-replay. Pseudo-predictions with confidence lower
216+
than this threshold are discarded for memory-replay
217+
218+
pytorch_cfg_updates: dict, optional, default = None.
219+
A dictionary of updates to the pytorch config. The keys are the dot-separated
220+
paths to the values to update in the config.
221+
For example, to update the gpus to run the training on, you can use:
222+
```
223+
pytorch_cfg_updates={"runner.gpus": [0,1,2,3]}
224+
```
195225
196226
Returns
197227
-------
@@ -255,20 +285,25 @@ def train_network(
255285
elif engine == Engine.PYTORCH:
256286
from deeplabcut.pose_estimation_pytorch.apis import train_network
257287

258-
_update_device(gputouse, torch_kwargs)
259-
if "display_iters" not in torch_kwargs:
260-
torch_kwargs["display_iters"] = displayiters
261-
262288
return train_network(
263289
config,
264290
shuffle=shuffle,
265291
trainingsetindex=trainingsetindex,
266292
modelprefix=modelprefix,
267-
max_snapshots_to_keep=max_snapshots_to_keep,
293+
device=device,
294+
snapshot_path=snapshot_path,
295+
detector_path=detector_path,
268296
load_head_weights=keepdeconvweights,
297+
batch_size=batch_size,
269298
epochs=epochs,
270299
save_epochs=save_epochs,
271-
**torch_kwargs,
300+
detector_batch_size=detector_batch_size,
301+
detector_epochs=detector_epochs,
302+
detector_save_epochs=detector_save_epochs,
303+
display_iters=displayiters,
304+
max_snapshots_to_keep=max_snapshots_to_keep,
305+
pose_threshold=pose_threshold,
306+
pytorch_cfg_updates=pytorch_cfg_updates,
272307
)
273308

274309
raise NotImplementedError(f"This function is not implemented for {engine}")

‎deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ def benchmark_paf_graphs(
143143

144144
# update the edges to keep in the PyTorch configuration file
145145
head_update = dict(predictor=dict(edges_to_keep=best_edges))
146-
loader.update_model_cfg(dict(model=dict(heads=dict(bodypart=head_update))))
146+
loader.update_model_cfg({"model.heads.bodypart": head_update})
147147

148148
# update the edges indices
149149
test_config = loader.model_folder.parent / "test" / "pose_cfg.yaml"

‎deeplabcut/pose_estimation_pytorch/apis/train.py‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ def train_network(
216216
display_iters: int | None = None,
217217
max_snapshots_to_keep: int | None = None,
218218
pose_threshold: float | None = 0.1,
219-
**kwargs,
219+
pytorch_cfg_updates: dict | None = None,
220220
) -> None:
221221
"""Trains a network for a project
222222
@@ -252,8 +252,14 @@ def train_network(
252252
max_snapshots_to_keep: the maximum number of snapshots to save for each model
253253
pose_threshold: Used for memory-replay. Pseudo-predictions with confidence lower
254254
than this threshold are discarded for memory-replay
255-
**kwargs : could be any entry of the pytorch_config dictionary. Examples are
256-
to see the full list see the pytorch_cfg.yaml file in your project folder
255+
pytorch_cfg_updates: dict, optional, default = None.
256+
A dictionary of updates to the pytorch config. The keys are the dot-separated
257+
paths to the values to update in the config.
258+
For example, to update the gpus to run the training on, you can use:
259+
```
260+
pytorch_cfg_updates={"runner.gpus": [0,1,2,3]}
261+
```
262+
To see the full list - check the pytorch_cfg.yaml file in your project folder
257263
"""
258264
loader = DLCLoader(
259265
config=config,
@@ -314,7 +320,9 @@ def train_network(
314320
if display_iters is not None:
315321
detector_cfg["train_settings"]["display_iters"] = display_iters
316322

317-
loader.update_model_cfg(kwargs)
323+
if pytorch_cfg_updates is not None:
324+
loader.update_model_cfg(pytorch_cfg_updates)
325+
318326
setup_file_logging(loader.model_folder / "train.txt")
319327

320328
logging.info("Training with configuration:")

‎deeplabcut/pose_estimation_pytorch/config/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,5 +19,6 @@
1919
pretty_print,
2020
read_config_as_dict,
2121
update_config,
22+
update_config_by_dotpath,
2223
write_config,
2324
)

‎deeplabcut/pose_estimation_pytorch/config/utils.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,47 @@ def update_config(config: dict, updates: dict, copy_original: bool = True) -> di
152152
return config
153153

154154

155+
def update_config_by_dotpath(config: dict, updates: dict, copy_original: bool = True) -> dict:
156+
"""Updates items in the configuration file using dot notation for nested keys
157+
158+
The configuration dict should only be composed of primitive Python types
159+
(dict, list and values). This is the case when reading the file using
160+
`read_config_as_dict`.
161+
162+
Args:
163+
config: the configuration dict to update
164+
updates: single-level dict with dot notation keys indicating nested paths
165+
e.g. {"device": "cuda", "runner.gpus": [0,1]}
166+
copy_original: whether to copy the original dict before updating it
167+
168+
Returns:
169+
the updated dictionary
170+
"""
171+
if copy_original:
172+
config = copy.deepcopy(config)
173+
174+
for key, value in updates.items():
175+
# Split key into parts by dots
176+
parts = key.split(".")
177+
178+
# Handle non-nested case
179+
if len(parts) == 1:
180+
config[key] = copy.deepcopy(value)
181+
continue
182+
183+
# Navigate to nested location
184+
current = config
185+
for part in parts[:-1]:
186+
if part not in current:
187+
current[part] = {}
188+
current = current[part]
189+
190+
# Set the value at final location
191+
current[parts[-1]] = copy.deepcopy(value)
192+
193+
return config
194+
195+
155196
def get_config_folder_path() -> Path:
156197
"""Returns: the Path to the folder containing the "configs" for DeepLabCut 3.0"""
157198
dlc_parent_path = Path(auxiliaryfunctions.get_deeplabcut_path())

‎deeplabcut/pose_estimation_pytorch/data/base.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def update_model_cfg(self, updates: dict) -> None:
6262
Args:
6363
updates: the items to update in the model configuration
6464
"""
65-
self.model_cfg = config.update_config(self.model_cfg, updates)
65+
self.model_cfg = config.update_config_by_dotpath(self.model_cfg, updates)
6666
config.write_config(self.model_config_path, self.model_cfg)
6767

6868
@abstractmethod

‎deeplabcut/pose_estimation_pytorch/modelzoo/train_from_coco.py‎

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -48,28 +48,21 @@ def adaptation_train(
4848

4949
utils.fix_seeds(loader.model_cfg["train_settings"]["seed"])
5050

51-
updates = dict(
52-
detector=dict(
53-
model=dict(freeze_bn_stats=True),
54-
runner=dict(snapshots=dict(max_snapshots=5, save_epochs=1)),
55-
train_settings=dict(batch_size=detector_batch_size, epochs=4),
56-
),
57-
model=dict(backbone=dict(freeze_bn_stats=True)),
58-
runner=dict(snapshots=dict(max_snapshots=5, save_epochs=1)),
59-
train_settings=dict(batch_size=batch_size, epochs=4),
60-
)
61-
62-
if epochs is not None:
63-
updates["train_settings"]["epochs"] = epochs
64-
if save_epochs is not None:
65-
updates["runner"]["snapshots"]["save_epochs"] = save_epochs
66-
if detector_epochs is not None:
67-
updates["detector"]["train_settings"]["epochs"] = detector_epochs
68-
if detector_save_epochs is not None:
69-
updates["detector"]["runner"]["snapshots"]["save_epochs"] = detector_save_epochs
51+
updates = {
52+
"detector.model.freeze_bn_stats": True,
53+
"detector.runner.snapshots.max_snapshots": 5,
54+
"detector.runner.snapshots.save_epochs": detector_save_epochs or 1,
55+
"detector.train_settings.batch_size": detector_batch_size,
56+
"detector.train_settings.epochs": detector_epochs or 4,
57+
"model.backbone.freeze_bn_stats": True,
58+
"runner.snapshots.max_snapshots": 5,
59+
"runner.snapshots.save_epochs": save_epochs or 1,
60+
"train_settings.batch_size": batch_size,
61+
"train_settings.epochs": epochs or 4,
62+
}
7063

7164
if eval_interval is not None:
72-
updates["runner"]["eval_interval"] = eval_interval
65+
updates["runner.eval_interval"] = eval_interval
7366

7467
loader.update_model_cfg(updates)
7568

‎tests/pose_estimation_pytorch/config/test_make_pose_config.py‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
make_basic_project_config,
1717
make_pytorch_pose_config,
1818
)
19-
from deeplabcut.pose_estimation_pytorch.config.utils import pretty_print, update_config
19+
from deeplabcut.pose_estimation_pytorch.config.utils import pretty_print, update_config, update_config_by_dotpath
2020

2121

2222
@pytest.mark.parametrize("bodyparts", [["nose"], ["nose", "ear", "eye"]])
@@ -398,6 +398,35 @@ def test_update_config(data: dict):
398398
assert result == data["expected_result"]
399399

400400

401+
@pytest.mark.parametrize("data", [
402+
{
403+
"config": {"a": 0, "b": 0},
404+
"updates": {"b": 1},
405+
"expected_result": {"a": 0, "b": 1},
406+
},
407+
{
408+
"config": {"a": 0, "b": {"i0": 1, "i1": 2}},
409+
"updates": {"b": 1},
410+
"expected_result": {"a": 0, "b": 1},
411+
},
412+
{
413+
"config": {"a": 0, "b": {"i0": 1, "i1": 2}},
414+
"updates": {"b.i0": [1, 2, 3]},
415+
"expected_result": {"a": 0, "b": {"i0": [1, 2, 3], "i1": 2}},
416+
},
417+
{
418+
"config": {"detector": {"batch_size": 1, "epochs": 10, "save_epochs": 5}},
419+
"updates": {"batch_size": 1, "detector.batch_size": 8, "detector.save_epochs": 1},
420+
"expected_result": {"batch_size": 1, "detector": {"batch_size": 8, "epochs": 10, "save_epochs": 1}},
421+
},
422+
])
423+
def test_update_config_by_dotpath(data: dict):
424+
result = update_config_by_dotpath(config=data["config"], updates=data["updates"])
425+
print("\nResult")
426+
pretty_print(result)
427+
assert result == data["expected_result"]
428+
429+
401430
def _make_project_config(
402431
project_path: str,
403432
multianimal: bool,

0 commit comments

Comments
 (0)