Function train_network

deeplabcut/pose_estimation_pytorch/apis/training.py:200–374

Trains a network for a project. Args: config : path to the yaml config file of the project shuffle : index of the shuffle we want to train on trainingsetindex : training set index modelprefix: directory containing the deeplabcut configuration files to use

(
    config: str | Path,
    shuffle: int = 1,
    trainingsetindex: int = 0,
    modelprefix: str = "",
    device: str | None = None,
    snapshot_path: str | Path | None = None,
    detector_path: str | Path | None = None,
    load_head_weights: bool = True,
    batch_size: int | None = None,
    epochs: int | None = None,
    save_epochs: int | None = None,
    detector_batch_size: int | None = None,
    detector_epochs: int | None = None,
    detector_save_epochs: int | None = None,
    display_iters: int | None = None,
    max_snapshots_to_keep: int | None = None,
    pose_threshold: float | None = 0.1,
    pytorch_cfg_updates: dict | None = None,
)

def train_network(
    config: str | Path,
    shuffle: int = 1,
    trainingsetindex: int = 0,
    modelprefix: str = "",
    device: str | None = None,
    snapshot_path: str | Path | None = None,
    detector_path: str | Path | None = None,
    load_head_weights: bool = True,
    batch_size: int | None = None,
    epochs: int | None = None,
    save_epochs: int | None = None,
    detector_batch_size: int | None = None,
    detector_epochs: int | None = None,
    detector_save_epochs: int | None = None,
    display_iters: int | None = None,
    max_snapshots_to_keep: int | None = None,
    pose_threshold: float | None = 0.1,
    pytorch_cfg_updates: dict | None = None,
) -> None:
    """Trains a network for a project.

    Args:
        config : path to the yaml config file of the project
        shuffle : index of the shuffle we want to train on
        trainingsetindex : training set index
        modelprefix: directory containing the deeplabcut configuration files to use
            to train the network (and where snapshots will be saved). By default, they
             are assumed to exist in the project folder.
        device: the torch device to train on (such as "cpu", "cuda", "mps")
        snapshot_path: if resuming training, the snapshot from which to resume
        detector_path: if resuming training of a top-down model, used to specify the
            detector snapshot from which to resume
        load_head_weights: if resuming training of a pose estimation model (either
            through the `snapshot_path` attribute or the `resume_training_from` key in
            the `pytorch_config.yaml` file), setting this to True also loads the weights
            for the model head (equivalent to the `keepdeconvweights` for  TensorFlow
            models). Note that if you change the number of bodyparts, you need to set
            this to false for re-training.
        batch_size: overrides the batch size to train with
        epochs: overrides the maximum number of epochs to train the model for
        save_epochs: overrides the number of epochs between each snapshot save
        detector_batch_size: Only for top-down models. Overrides the batch size with
            which to train the detector.
        detector_epochs: Only for top-down models. Overrides the maximum number of
            epochs to train the model for. Setting to 0 means the detector will not be
            trained.
        detector_save_epochs: Only for top-down models. Overrides the number of epochs
            between each snapshot of the detector is saved.
        display_iters: overrides the number of iterations between each log of the loss
            within an epoch
        max_snapshots_to_keep: the maximum number of snapshots to save for each model
        pose_threshold: Used for memory-replay. Pseudo-predictions with confidence lower
            than this threshold are discarded for memory-replay
        pytorch_cfg_updates: dict, optional, default = None.
            A dictionary of updates to the pytorch config. The keys are the dot-separated
            paths to the values to update in the config.
            For example, to update the gpus to run the training on, you can use: