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:

Name	Type	Description	Default
config	string	Full path of the config.yaml file as a string.	required
shuffle	int	Integer value specifying the shuffle index to select for training.	1
trainingsetindex	int	Which TrainingsetFraction to use. By default the first (TrainingFraction is a list in config.yaml). Defaults to 0.	0

Returns:

Name	Type	Description
tuple		trainposeconfigfile, testposeconfigfile, snapshotfolder.

Source code in deeplabcut/pose_estimation_tensorflow/training.py

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.

    Args:
        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): Which TrainingsetFraction to use. By default the first
            (TrainingFraction is a list in config.yaml). Defaults to 0.

    Returns:
        tuple: 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(cfg["project_path"]) / str(modelfoldername) / "train" / "pose_cfg.yaml"
    testposeconfigfile = Path(cfg["project_path"]) / str(modelfoldername) / "test" / "pose_cfg.yaml"
    snapshotfolder = Path(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:

Name	Type	Description	Default
config	string	Full path of the config.yaml file as a string.	required
shuffle	int	Integer value specifying the shuffle index to select for training. Defaults to 1.	1
trainingsetindex	int	Integer specifying which TrainingsetFraction to use. Note that TrainingFraction is a list in config.yaml. Defaults to 0.	0
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	5
display_iters	int	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! Defaults to None.	required
saveiters	int	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! Defaults to None.	None
maxiters	int	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! Defaults to None.	None
allow_growth	bool	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. Defaults to True.	True
gputouse	int	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. Defaults to None.	None
autotune	bool	Property of TensorFlow, somehow faster if False (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317). Defaults to False.	False
keepdeconvweights	bool	Restores deconvolution layer (and backbone) weights when training from a snapshot. Set to false if bodypart count changes. Defaults to True.	True
modelprefix	str	Directory containing the deeplabcut models to use when evaluating the network. By default, the models are assumed to exist in the project folder. Defaults to "".	''
superanimal_name	str	Specified if transfer learning with superanimal is desired. Defaults to "".	''
superanimal_transfer_learning	bool	If true, transfer learning (new decoding layer). If false and superanimal_name is set, fine-tuning (reuse decoding layer). Defaults to False.	False

Returns:

Type	Description
	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

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.

    Args:
        config (string): Full path of the config.yaml file as a string.
        shuffle (int, optional): Integer value specifying the shuffle index to select for training. Defaults to 1.
        trainingsetindex (int, optional): Integer specifying which TrainingsetFraction to use.
            Note that TrainingFraction is a list in config.yaml. Defaults to 0.
        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
        display_iters (int, optional): 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! Defaults to None.
        saveiters (int, optional): 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! Defaults to None.
        maxiters (int, optional): 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! Defaults to None.
        allow_growth (bool, optional): 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. Defaults to True.
        gputouse (int, optional): 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. Defaults to None.
        autotune (bool, optional): Property of TensorFlow, somehow faster if ``False``
            (as Eldar found out, see https://github.com/tensorflow/tensorflow/issues/13317). Defaults to False.
        keepdeconvweights (bool, optional): Restores deconvolution layer (and backbone) weights when
            training from a snapshot. Set to false if bodypart count changes. Defaults to True.
        modelprefix (str, optional): Directory containing the deeplabcut models to use when evaluating the network.
            By default, the models are assumed to exist in the project folder. Defaults to "".
        superanimal_name (str, optional): Specified if transfer learning with superanimal is desired. Defaults to "".
        superanimal_transfer_learning (bool, optional): If true, transfer learning (new decoding layer).
            If false and superanimal_name is set, fine-tuning (reuse decoding layer). Defaults to False.

    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 = Path.cwd()

    # 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(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 != "":
            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 Path(weight_folder).exists():
                    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 = list(Path(weight_folder).glob("snapshot-*.index"))
            init_weights = auxiliaryfunctions.safe_resolve(snapshots[0].with_suffix(""))

            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=str(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."
    )