Visitar URL original
Dynamic update of pytorch config with dot notation (#2859) · DeepLabCut/DeepLabCut@45d4d3d · GitHub
Skip to content

Commit 45d4d3d

Browse files
maximpavlivmaximpavliv
authored
Dynamic update of pytorch config with dot notation (#2859)
1 parent e128877 commit 45d4d3d

11 files changed

Lines changed: 267 additions & 140 deletions

File tree

‎deeplabcut/compat.py‎

Lines changed: 55 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,15 @@ 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
):
8189
"""
8290
Trains the network with the labels in the training dataset.
@@ -135,6 +143,7 @@ def train_network(
135143
Only for the PyTorch engine (equivalent to the `saveiters` parameter for the
136144
TensorFlow engine). The number of epochs between each snapshot save. If
137145
None, the value will be read from the `pytorch_config.yaml` file.
146+
138147
allow_growth: bool, optional, default=True.
139148
Only for the TensorFlow engine.
140149
For some smaller GPUs the memory issues happen. If ``True``, the memory
@@ -180,18 +189,38 @@ def train_network(
180189
overwrite this by passing the engine as an argument, but this should generally
181190
not be done.
182191
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)
192+
device: str, optional, default = None.
193+
Only for the PyTorch engine. The device to run the training on (e.g. "cuda:0")
194+
195+
snapshot_path: str or Path, optional, default = None.
196+
Only for the PyTorch engine. The path to the pose model snapshot to resume training from.
197+
198+
detector_path: str or Path, optional, default = None.
199+
Only for the PyTorch engine. The path to the detector model snapshot to resume training from.
200+
201+
batch_size: int, optional, default = None.
202+
Only for the PyTorch engine. The batch size to use while training.
203+
204+
detector_batch_size: int, optional, default = None.
205+
Only for the PyTorch engine. The batch size to use while training the detector.
206+
207+
detector_epochs: int, optional, default = None.
208+
Only for the PyTorch engine. The number of epochs to train the detector for.
209+
210+
detector_save_epochs: int, optional, default = None.
211+
Only for the PyTorch engine. The number of epochs between each detector snapshot save.
191212
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``.
213+
pose_threshold: float, optional, default = 0.1.
214+
Only for the PyTorch engine. Used for memory-replay. Pseudo-predictions with confidence lower
215+
than this threshold are discarded for memory-replay
216+
217+
pytorch_cfg_updates: dict, optional, default = None.
218+
A dictionary of updates to the pytorch config. The keys are the dot-separated
219+
paths to the values to update in the config.
220+
For example, to update the gpus to run the training on, you can use:
221+
```
222+
pytorch_cfg_updates={"runner.gpus": [0,1,2,3]}
223+
```
195224
196225
Returns
197226
-------
@@ -255,20 +284,25 @@ def train_network(
255284
elif engine == Engine.PYTORCH:
256285
from deeplabcut.pose_estimation_pytorch.apis import train_network
257286

258-
_update_device(gputouse, torch_kwargs)
259-
if "display_iters" not in torch_kwargs:
260-
torch_kwargs["display_iters"] = displayiters
261-
262287
return train_network(
263288
config,
264289
shuffle=shuffle,
265290
trainingsetindex=trainingsetindex,
266291
modelprefix=modelprefix,
267-
max_snapshots_to_keep=max_snapshots_to_keep,
292+
device=device,
293+
snapshot_path=snapshot_path,
294+
detector_path=detector_path,
268295
load_head_weights=keepdeconvweights,
296+
batch_size=batch_size,
269297
epochs=epochs,
270298
save_epochs=save_epochs,
271-
**torch_kwargs,
299+
detector_batch_size=detector_batch_size,
300+
detector_epochs=detector_epochs,
301+
detector_save_epochs=detector_save_epochs,
302+
display_iters=displayiters,
303+
max_snapshots_to_keep=max_snapshots_to_keep,
304+
pose_threshold=pose_threshold,
305+
pytorch_cfg_updates=pytorch_cfg_updates,
272306
)
273307

274308
raise NotImplementedError(f"This function is not implemented for {engine}")
@@ -1110,6 +1144,7 @@ def analyze_images(
11101144

11111145
if engine == Engine.PYTORCH:
11121146
from deeplabcut.pose_estimation_pytorch import analyze_images
1147+
11131148
return analyze_images(
11141149
config=config,
11151150
images=images,
@@ -1218,6 +1253,7 @@ def analyze_time_lapse_frames(
12181253

12191254
if engine == Engine.TF:
12201255
from deeplabcut.pose_estimation_tensorflow import analyze_time_lapse_frames
1256+
12211257
return analyze_time_lapse_frames(
12221258
config,
12231259
directory,
@@ -1230,6 +1266,7 @@ def analyze_time_lapse_frames(
12301266
)
12311267
elif engine == Engine.PYTORCH:
12321268
from deeplabcut.pose_estimation_pytorch import analyze_images
1269+
12331270
return analyze_images(
12341271
config=config,
12351272
images=directory,

‎deeplabcut/pose_estimation_pytorch/apis/prune_paf_graph.py‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,12 @@ def benchmark_paf_graphs(
7878

7979
gt_train = loader.ground_truth_keypoints("train")
8080
best_paf_edges, _ = get_n_best_paf_graphs(
81-
model, gt_train, preprocessor, device, predictor.graph, n_graphs=10,
81+
model,
82+
gt_train,
83+
preprocessor,
84+
device,
85+
predictor.graph,
86+
n_graphs=10,
8287
)
8388

8489
if verbose:
@@ -142,14 +147,13 @@ def benchmark_paf_graphs(
142147
print()
143148

144149
# update the edges to keep in the PyTorch configuration file
145-
head_update = dict(predictor=dict(edges_to_keep=best_edges))
146-
loader.update_model_cfg(dict(model=dict(heads=dict(bodypart=head_update))))
150+
loader.update_model_cfg(
151+
{"model.heads.bodypart.predictor.edges_to_keep": best_edges}
152+
)
147153

148154
# update the edges indices
149155
test_config = loader.model_folder.parent / "test" / "pose_cfg.yaml"
150-
auxiliaryfunctions.edit_config(
151-
str(test_config), dict(paf_best=best_edges)
152-
)
156+
auxiliaryfunctions.edit_config(str(test_config), dict(paf_best=best_edges))
153157

154158
return results
155159

‎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: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ def replace_default_values(
5959
ValueError: if there is a placeholder value who's "updated" value was not
6060
given to the method
6161
"""
62+
6263
def get_updated_value(variable: str) -> int | list[int]:
6364
var_parts = variable.strip().split(" ")
6465
var_name = var_parts[0]
@@ -152,6 +153,49 @@ def update_config(config: dict, updates: dict, copy_original: bool = True) -> di
152153
return config
153154

154155

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

207251
return cfg
208252

‎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

‎examples/testscript_pytorch_multi_animal.py‎

Lines changed: 20 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
#
99
# Licensed under GNU Lesser General Public License v3.0
1010
#
11-
""" Testscript for single animal PyTorch projects """
11+
"""Testscript for single animal PyTorch projects"""
1212
from __future__ import annotations
1313

1414
from pathlib import Path
@@ -17,7 +17,11 @@
1717
from deeplabcut.compat import Engine
1818

1919
from utils import (
20-
cleanup, create_fake_project, log_step, run, SyntheticProjectParameters,
20+
cleanup,
21+
create_fake_project,
22+
log_step,
23+
run,
24+
SyntheticProjectParameters,
2125
)
2226

2327

@@ -57,35 +61,21 @@ def main(
5761
net_type=net_type,
5862
videos=[str(project_path / "videos" / "video.mp4")],
5963
device=device,
60-
train_kwargs=dict(
61-
train_settings=dict(
62-
display_iters=50,
63-
epochs=epochs_,
64-
batch_size=batch_size,
65-
),
66-
runner=dict(
67-
device=device,
68-
snapshots=dict(
69-
save_epochs=save_epochs,
70-
max_snapshots=max_snapshots_to_keep,
71-
)
72-
),
73-
detector=dict(
74-
train_settings=dict(
75-
display_iters=1,
76-
epochs=detector_epochs,
77-
batch_size=detector_batch_size,
78-
),
79-
runner=dict(
80-
snapshots=dict(
81-
save_epochs=save_epochs,
82-
max_snapshots=max_snapshots_to_keep,
83-
)
84-
)
85-
),
86-
logger=logger,
87-
),
8864
engine=engine,
65+
pytorch_cfg_updates={
66+
"train_settings.display_iters": 50,
67+
"train_settings.epochs": epochs_,
68+
"train_settings.batch_size": batch_size,
69+
"runner.device": device,
70+
"runner.snapshots.save_epochs": save_epochs,
71+
"runner.snapshots.max_snapshots": max_snapshots_to_keep,
72+
"detector.train_settings.display_iters": 1,
73+
"detector.train_settings.epochs": detector_epochs,
74+
"detector.train_settings.batch_size": detector_batch_size,
75+
"detector.runner.snapshots.save_epochs": save_epochs,
76+
"detector.runner.snapshots.max_snapshots": max_snapshots_to_keep,
77+
"logger": logger,
78+
},
8979
create_labeled_videos=create_labeled_videos,
9080
)
9181
except Exception as err:

0 commit comments

Comments
 (0)