Skip to content

deeplabcut.pose_estimation_tensorflow.training

Functions:

Name Description
return_train_network_path

Returns the training and test pose config file names as well as the folder where
train_network

Trains the network with the labels in the training dataset.

return_train_network_path 

return_train_network_path(config, shuffle=1, trainingsetindex=0, modelprefix='')

Returns the training and test pose config file names as well as the folder where the snapshot is.

Parameters

config : string Full path of the config.yaml file as a string.

int

Integer value specifying the shuffle index to select for training.

int, optional

Integer specifying which TrainingsetFraction to use. By default the first (note that TrainingFraction is a list in config.yaml).

Returns the triple: trainposeconfigfile, testposeconfigfile, snapshotfolder

Source code in deeplabcut/pose_estimation_tensorflow/training.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def return_train_network_path(config, shuffle=1, trainingsetindex=0, modelprefix=""):
    """Returns the training and test pose config file names as well as the folder where
    the snapshot is.

    Parameters
    ----------
    config : string
        Full path of the config.yaml file as a string.

    shuffle: int
        Integer value specifying the shuffle index to select for training.

    trainingsetindex: int, optional
        Integer specifying which TrainingsetFraction to use. By default the first (note that TrainingFraction is a list
        in config.yaml).

    Returns the triple: trainposeconfigfile, testposeconfigfile, snapshotfolder
    """
    from deeplabcut.utils import auxiliaryfunctions

    cfg = auxiliaryfunctions.read_config(config)
    modelfoldername = auxiliaryfunctions.get_model_folder(
        cfg["TrainingFraction"][trainingsetindex], shuffle, cfg, modelprefix=modelprefix
    )
    trainposeconfigfile = Path(os.path.join(cfg["project_path"], str(modelfoldername), "train", "pose_cfg.yaml"))
    testposeconfigfile = Path(os.path.join(cfg["project_path"], str(modelfoldername), "test", "pose_cfg.yaml"))
    snapshotfolder = Path(os.path.join(cfg["project_path"], str(modelfoldername), "train"))

    return trainposeconfigfile, testposeconfigfile, snapshotfolder

train_network 

train_network(
    config,
    shuffle=1,
    trainingsetindex=0,
    max_snapshots_to_keep=5,
    displayiters=None,
    saveiters=None,
    maxiters=None,
    allow_growth=True,
    gputouse=None,
    autotune=False,
    keepdeconvweights=True,
    modelprefix="",
    superanimal_name="",
    superanimal_transfer_learning=False,
)

Trains the network with the labels in the training dataset.

Parameters
----------
config : string
    Full path of the config.yaml file as a string.

shuffle: int, optional, default=1
    Integer value specifying the shuffle index to select for training.

trainingsetindex: int, optional, default=0
    Integer specifying which TrainingsetFraction to use.
    Note that TrainingFraction is a list in config.yaml.

max_snapshots_to_keep: int or None
    Sets how many snapshots are kept, i.e. states of the trained network. Every
    saving iteration many times a snapshot is stored, however only the last
    ``max_snapshots_to_keep`` many are kept! If you change this to None, then all
    are kept.
    See: https://github.com/DeepLabCut/DeepLabCut/issues/8#issuecomment-387404835

displayiters: optional, default=None
    This variable is actually set in ``pose_config.yaml``. However, you can
    overwrite it with this hack. Don't use this regularly, just if you are too lazy
    to dig out the ``pose_config.yaml`` file for the corresponding project. If
    ``None``, the value from there is used, otherwise it is overwritten!

saveiters: optional, default=None
    This variable is actually set in ``pose_config.yaml``. However, you can
    overwrite it with this hack. Don't use this regularly, just if you are too lazy
    to dig out the ``pose_config.yaml`` file for the corresponding project.
    If ``None``, the value from there is used, otherwise it is overwritten!

maxiters: optional, default=None
    This variable is actually set in ``pose_config.yaml``. However, you can
    overwrite it with this hack. Don't use this regularly, just if you are too lazy
    to dig out the ``pose_config.yaml`` file for the corresponding project.
    If ``None``, the value from there is used, otherwise it is overwritten!

allow_growth: bool, optional, default=True.
    For some smaller GPUs the memory issues happen. If ``True``, the memory
    allocator does not pre-allocate the entire specified GPU memory region, instead
    starting small and growing as needed.
    See issue: https://forum.image.sc/t/how-to-stop-running-out-of-vram/30551/2

gputouse: optional, default=None
    Natural number indicating the number of your GPU (see number in nvidia-smi).
    If you do not have a GPU put None.
    See: https://nvidia.custhelp.com/app/answers/detail/a_id/3751/~/useful-nvidia-smi-queries

autotune: bool, optional, default=False
    Property of TensorFlow, somehow faster if ``False``
    (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317).

keepdeconvweights: bool, optional, default=True
    Also restores the weights of the deconvolution layers (and the backbone) when
    training from a snapshot. Note that if you change the number of bodyparts, you
    need to set this to false for re-training.

modelprefix: str, optional, default=""
    Directory containing the deeplabcut models to use when evaluating the network.
    By default, the models are assumed to exist in the project folder.

superanimal_name: str, optional, default =""
    Specified if transfer learning with superanimal is desired

superanimal_transfer_learning: bool, optional, default = False.
    If set true, the training is transfer learning (new decoding layer). If set false,

and superanimal_name is True, then the training is fine-tuning (reusing the decoding layer)

Returns
-------
None

Examples
--------
To train the network for first shuffle of the training dataset

>>> deeplabcut.train_network('/analysis/project/reaching-task/config.yaml')

To train the network for second shuffle of the training dataset

>>> deeplabcut.train_network(
        '/analysis/project/reaching-task/config.yaml',
        shuffle=2,
        keepdeconvweights=True,
    )
Source code in deeplabcut/pose_estimation_tensorflow/training.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
def train_network(
    config,
    shuffle=1,
    trainingsetindex=0,
    max_snapshots_to_keep=5,
    displayiters=None,
    saveiters=None,
    maxiters=None,
    allow_growth=True,
    gputouse=None,
    autotune=False,
    keepdeconvweights=True,
    modelprefix="",
    superanimal_name="",
    superanimal_transfer_learning=False,
):
    """Trains the network with the labels in the training dataset.

        Parameters
        ----------
        config : string
            Full path of the config.yaml file as a string.

        shuffle: int, optional, default=1
            Integer value specifying the shuffle index to select for training.

        trainingsetindex: int, optional, default=0
            Integer specifying which TrainingsetFraction to use.
            Note that TrainingFraction is a list in config.yaml.

        max_snapshots_to_keep: int or None
            Sets how many snapshots are kept, i.e. states of the trained network. Every
            saving iteration many times a snapshot is stored, however only the last
            ``max_snapshots_to_keep`` many are kept! If you change this to None, then all
            are kept.
            See: https://github.com/DeepLabCut/DeepLabCut/issues/8#issuecomment-387404835

        displayiters: optional, default=None
            This variable is actually set in ``pose_config.yaml``. However, you can
            overwrite it with this hack. Don't use this regularly, just if you are too lazy
            to dig out the ``pose_config.yaml`` file for the corresponding project. If
            ``None``, the value from there is used, otherwise it is overwritten!

        saveiters: optional, default=None
            This variable is actually set in ``pose_config.yaml``. However, you can
            overwrite it with this hack. Don't use this regularly, just if you are too lazy
            to dig out the ``pose_config.yaml`` file for the corresponding project.
            If ``None``, the value from there is used, otherwise it is overwritten!

        maxiters: optional, default=None
            This variable is actually set in ``pose_config.yaml``. However, you can
            overwrite it with this hack. Don't use this regularly, just if you are too lazy
            to dig out the ``pose_config.yaml`` file for the corresponding project.
            If ``None``, the value from there is used, otherwise it is overwritten!

        allow_growth: bool, optional, default=True.
            For some smaller GPUs the memory issues happen. If ``True``, the memory
            allocator does not pre-allocate the entire specified GPU memory region, instead
            starting small and growing as needed.
            See issue: https://forum.image.sc/t/how-to-stop-running-out-of-vram/30551/2

        gputouse: optional, default=None
            Natural number indicating the number of your GPU (see number in nvidia-smi).
            If you do not have a GPU put None.
            See: https://nvidia.custhelp.com/app/answers/detail/a_id/3751/~/useful-nvidia-smi-queries

        autotune: bool, optional, default=False
            Property of TensorFlow, somehow faster if ``False``
            (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317).

        keepdeconvweights: bool, optional, default=True
            Also restores the weights of the deconvolution layers (and the backbone) when
            training from a snapshot. Note that if you change the number of bodyparts, you
            need to set this to false for re-training.

        modelprefix: str, optional, default=""
            Directory containing the deeplabcut models to use when evaluating the network.
            By default, the models are assumed to exist in the project folder.

        superanimal_name: str, optional, default =""
            Specified if transfer learning with superanimal is desired

        superanimal_transfer_learning: bool, optional, default = False.
            If set true, the training is transfer learning (new decoding layer). If set false,
    and superanimal_name is True, then the training is fine-tuning (reusing the decoding layer)

        Returns
        -------
        None

        Examples
        --------
        To train the network for first shuffle of the training dataset

        >>> deeplabcut.train_network('/analysis/project/reaching-task/config.yaml')

        To train the network for second shuffle of the training dataset

        >>> deeplabcut.train_network(
                '/analysis/project/reaching-task/config.yaml',
                shuffle=2,
                keepdeconvweights=True,
            )
    """
    if allow_growth:
        os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"

    # reload logger.
    import importlib
    import logging

    import tensorflow as tf

    importlib.reload(logging)
    logging.shutdown()

    from deeplabcut.utils import auxiliaryfunctions

    tf.compat.v1.reset_default_graph()
    start_path = os.getcwd()

    # Read file path for pose_config file. >> pass it on
    cfg = auxiliaryfunctions.read_config(config)
    modelfoldername = auxiliaryfunctions.get_model_folder(
        cfg["TrainingFraction"][trainingsetindex], shuffle, cfg, modelprefix=modelprefix
    )
    poseconfigfile = Path(os.path.join(cfg["project_path"], str(modelfoldername), "train", "pose_cfg.yaml"))
    if not poseconfigfile.is_file():
        print("The training datafile ", poseconfigfile, " is not present.")
        print("Probably, the training dataset for this specific shuffle index was not created.")
        print(
            "Try with a different shuffle/trainingsetfraction or use function 'create_training_dataset' to create a new"
            "trainingdataset with this shuffle index."
        )
    else:
        # Set environment variables
        if autotune is not False:  # see: https://github.com/tensorflow/tensorflow/issues/13317
            os.environ["TF_CUDNN_USE_AUTOTUNE"] = "0"
        if gputouse is not None:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(gputouse)
    try:
        cfg_dlc = auxiliaryfunctions.read_plainconfig(poseconfigfile)

        if superanimal_name != "":
            import glob

            from dlclibrary.dlcmodelzoo.modelzoo_download import (
                MODELOPTIONS,
                download_huggingface_model,
            )

            from deeplabcut.modelzoo.utils import parse_available_supermodels

            dlc_root_path = auxiliaryfunctions.get_deeplabcut_path()
            parse_available_supermodels()
            weight_folder = str(
                Path(dlc_root_path)
                / "pose_estimation_tensorflow"
                / "models"
                / "pretrained"
                / (superanimal_name + "_weights")
            )

            if superanimal_name in MODELOPTIONS:
                if not os.path.exists(weight_folder):
                    download_huggingface_model(superanimal_name, weight_folder)
                else:
                    print(f"{weight_folder} exists, using the downloaded weights")
            else:
                print(
                    f"{superanimal_name} not available. Available ones are: ",
                    MODELOPTIONS,
                )

            snapshots = glob.glob(os.path.join(weight_folder, "snapshot-*.index"))
            init_weights = os.path.abspath(snapshots[0]).replace(".index", "")

            from deeplabcut.pose_estimation_tensorflow.core.train_multianimal import (
                train,
            )

            print("Selecting multi-animal trainer")
            train(
                str(poseconfigfile),
                displayiters,
                saveiters,
                maxiters,
                max_to_keep=max_snapshots_to_keep,
                keepdeconvweights=keepdeconvweights,
                allow_growth=allow_growth,
                init_weights=init_weights,
                remove_head=(True if superanimal_name != "" and superanimal_transfer_learning else False),
            )  # pass on path and file name for pose_cfg.yaml!

        elif "multi-animal" in cfg_dlc["dataset_type"]:
            from deeplabcut.pose_estimation_tensorflow.core.train_multianimal import (
                train,
            )

            print("Selecting multi-animal trainer")
            train(
                str(poseconfigfile),
                displayiters,
                saveiters,
                maxiters,
                max_to_keep=max_snapshots_to_keep,
                keepdeconvweights=keepdeconvweights,
                allow_growth=allow_growth,
            )  # pass on path and file name for pose_cfg.yaml!
        else:
            from deeplabcut.pose_estimation_tensorflow.core.train import train

            print("Selecting single-animal trainer")
            train(
                str(poseconfigfile),
                displayiters,
                saveiters,
                maxiters,
                max_to_keep=max_snapshots_to_keep,
                keepdeconvweights=keepdeconvweights,
                allow_growth=allow_growth,
            )  # pass on path and file name for pose_cfg.yaml!

    except BaseException as e:
        raise e
    finally:
        os.chdir(str(start_path))
    print(
        "The network is now trained and ready to evaluate. Use the function 'evaluate_network' to evaluate the network."
    )