deeplabcut.pose_estimation_pytorch.runners.base

Classes:

Name	Description
Runner	Runner base class.

Functions:

Name	Description
attempt_snapshot_load	Attempts to load a snapshot using torch.load(...).
fix_snapshot_metadata	Replace numpy floats in snapshot metrics.
get_load_weights_only	Gets the default value to use when loading snapshots with torch.load(...).
set_load_weights_only	Sets the default value to use when loading snapshots with torch.load(...).

Runner

Bases: ABC, Generic[ModelType]

Runner base class.

A runner takes a model and runs actions on it, such as training or inference

Methods:

Name	Description
__init__	Args:
load_snapshot	Loads the state dict for a model from a file.

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

class Runner(ABC, Generic[ModelType]):
    """Runner base class.

    A runner takes a model and runs actions on it, such as training or inference
    """

    def __init__(
        self,
        model: ModelType,
        device: str = "cpu",
        gpus: list[int] | None = None,
        snapshot_path: str | Path | None = None,
    ):
        """
        Args:
            model: the model to run
            device: the device to use (e.g. {'cpu', 'cuda:0', 'mps'})
            gpus: the list of GPU indices to use for multi-GPU training
            snapshot_path: the path of a snapshot from which to load model weights
        """
        if gpus is None:
            gpus = []

        if len(gpus) == 1:
            if device != "cuda":
                raise ValueError(
                    f"When specifying a GPU index to train on, the device must be set to 'cuda'. Found {device}"
                )
            device = f"cuda:{gpus[0]}"

        self.model = model
        self.device = device
        self.snapshot_path = snapshot_path
        self._gpus = gpus
        self._data_parallel = len(gpus) > 1

    @staticmethod
    def load_snapshot(
        snapshot_path: str | Path,
        device: str,
        model: ModelType,
        weights_only: bool | None = None,
    ) -> dict:
        """Loads the state dict for a model from a file.

        This method loads a file containing a DeepLabCut PyTorch model snapshot onto
        a given device, and sets the model weights using the state_dict.

        Args:
            snapshot_path: The path containing the model weights to load
            device: The device on which the model should be loaded
            model: The model for which the weights are loaded
            weights_only: Value for torch.load() `weights_only` parameter.
                If False, the python pickle module is used implicitly, which is known to
                be insecure. Only set to False if you're loading data that you trust
                (e.g. snapshots that you created yourself). For more information, see:
                    https://pytorch.org/docs/stable/generated/torch.load.html
                If None, the default value is used:
                    `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`

        Returns:
            The content of the snapshot file.
        """
        snapshot = attempt_snapshot_load(snapshot_path, device, weights_only)
        model.load_state_dict(snapshot["model"])
        return snapshot

__init__

__init__(model: ModelType, device: str = 'cpu', gpus: list[int] | None = None, snapshot_path: str | Path | None = None)

Parameters:

Name	Type	Description	Default
model	ModelType	the model to run	required
device	str	the device to use (e.g. {'cpu', 'cuda:0', 'mps'})	'cpu'
gpus	list[int] | None	the list of GPU indices to use for multi-GPU training	None
snapshot_path	str | Path | None	the path of a snapshot from which to load model weights	None

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

def __init__(
    self,
    model: ModelType,
    device: str = "cpu",
    gpus: list[int] | None = None,
    snapshot_path: str | Path | None = None,
):
    """
    Args:
        model: the model to run
        device: the device to use (e.g. {'cpu', 'cuda:0', 'mps'})
        gpus: the list of GPU indices to use for multi-GPU training
        snapshot_path: the path of a snapshot from which to load model weights
    """
    if gpus is None:
        gpus = []

    if len(gpus) == 1:
        if device != "cuda":
            raise ValueError(
                f"When specifying a GPU index to train on, the device must be set to 'cuda'. Found {device}"
            )
        device = f"cuda:{gpus[0]}"

    self.model = model
    self.device = device
    self.snapshot_path = snapshot_path
    self._gpus = gpus
    self._data_parallel = len(gpus) > 1

load_snapshot staticmethod

load_snapshot(snapshot_path: str | Path, device: str, model: ModelType, weights_only: bool | None = None) -> dict

Loads the state dict for a model from a file.

This method loads a file containing a DeepLabCut PyTorch model snapshot onto a given device, and sets the model weights using the state_dict.

Parameters:

Name	Type	Description	Default
snapshot_path	str | Path	The path containing the model weights to load	required
device	str	The device on which the model should be loaded	required
model	ModelType	The model for which the weights are loaded	required
weights_only	bool | None	Value for torch.load() weights_only parameter. If False, the python pickle module is used implicitly, which is known to be insecure. Only set to False if you're loading data that you trust (e.g. snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html If None, the default value is used: deeplabcut.pose_estimation_pytorch.get_load_weights_only()	None

Returns:

Type	Description
dict	The content of the snapshot file.

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

@staticmethod
def load_snapshot(
    snapshot_path: str | Path,
    device: str,
    model: ModelType,
    weights_only: bool | None = None,
) -> dict:
    """Loads the state dict for a model from a file.

    This method loads a file containing a DeepLabCut PyTorch model snapshot onto
    a given device, and sets the model weights using the state_dict.

    Args:
        snapshot_path: The path containing the model weights to load
        device: The device on which the model should be loaded
        model: The model for which the weights are loaded
        weights_only: Value for torch.load() `weights_only` parameter.
            If False, the python pickle module is used implicitly, which is known to
            be insecure. Only set to False if you're loading data that you trust
            (e.g. snapshots that you created yourself). For more information, see:
                https://pytorch.org/docs/stable/generated/torch.load.html
            If None, the default value is used:
                `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`

    Returns:
        The content of the snapshot file.
    """
    snapshot = attempt_snapshot_load(snapshot_path, device, weights_only)
    model.load_state_dict(snapshot["model"])
    return snapshot

attempt_snapshot_load

attempt_snapshot_load(path: str | Path, device: str, weights_only: bool | None = None) -> dict

Attempts to load a snapshot using torch.load(...).

Parameters:

Name	Type	Description	Default
path	str | Path	The path of the snapshot to try to load..	required
device	str	The device to use for the map_location.	required
weights_only	bool | None	Value for torch.load() weights_only parameter. If False, the python pickle module is used implicitly, which is known to be insecure. Only set to False if you're loading data that you trust (e.g. snapshots that you created yourself). For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html If None, the default value is used: deeplabcut.pose_estimation_pytorch.get_load_weights_only()	None

Returns:

Type	Description
dict	The loaded snapshot.

Raises:

Type	Description
UnpicklingError	If weights_only=True but the snapshot failed to load with weights_only=True.

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

def attempt_snapshot_load(
    path: str | Path,
    device: str,
    weights_only: bool | None = None,
) -> dict:
    """Attempts to load a snapshot using `torch.load(...)`.

    Args:
        path: The path of the snapshot to try to load..
        device: The device to use for the `map_location`.
        weights_only: Value for torch.load() `weights_only` parameter.
            If False, the python pickle module is used implicitly, which is known to be
            insecure. Only set to False if you're loading data that you trust (e.g.
            snapshots that you created yourself). For more information, see:
                https://pytorch.org/docs/stable/generated/torch.load.html
            If None, the default value is used:
                `deeplabcut.pose_estimation_pytorch.get_load_weights_only()`

    Returns:
        The loaded snapshot.

    Raises:
        pickle.UnpicklingError: If `weights_only=True` but the snapshot failed to load
            with `weights_only=True`.
    """
    try:
        if weights_only is None:
            weights_only = get_load_weights_only()

        snapshot = torch.load(path, map_location=device, weights_only=weights_only)
    except pickle.UnpicklingError as err:
        logging.error(
            f"\nFailed to load the snapshot: {path}.\n\n"
            "If you trust the snapshot that you're trying to load, you can try\n"
            "calling `Runner.load_snapshot` with `weights_only=False`. See the \n"
            "error message below for more information and warnings.\n"
            "You can set the `weights_only` parameter in the model configuration (\n"
            "the content of the pytorch_config.yaml), as:\n\n```\n"
            "runner:\n"
            "  load_weights_only: False\n```\n\n"
            "If it's the detector snapshot that's failing to load, place the\n"
            "`load_weights_only` key under the detector runner:\n\n```\n"
            "detector:\n"
            "    runner:\n"
            "      load_weights_only: False\n```\n\n"
            "You can also set the default `load_weights_only` that will be used when\n"
            "the `load_weights_only` variable is not set in the `pytorch_config.yaml`\n"
            "using `deeplabcut.pose_estimation_pytorch.set_load_weights_only(value)`:\n"
            "\n```\n"
            "from deeplabcut.pose_estimation_pytorch import set_load_weights_only\n"
            "set_load_weights_only(True)\n"
            "```\n\n"
            "You can also set the value for `load_weights_only` with a \n"
            "`TORCH_LOAD_WEIGHTS_ONLY` environment variable. If you call \n"
            "`TORCH_LOAD_WEIGHTS_ONLY=False python -m deeplabcut`, it will launch the\n"
            "DeepLabCut GUI with the default `load_weights_only` value to False.\n"
            "If you set this value to `False`, make sure you only load snapshots that\n"
            "you trust.\n\n"
        )
        raise err

    return snapshot

fix_snapshot_metadata

fix_snapshot_metadata(path: str | Path) -> None

Replace numpy floats in snapshot metrics.

Only call this method with snapshots that you trust, as torch.load(...) is called with weights_only=False. For more information, see: https://pytorch.org/docs/stable/generated/torch.load.html

DeepLabCut PyTorch snapshots trained with older releases may have numpy floats in the stored metrics. This method opens the snapshots (with weights_only=False), replaces the numpy floats with python floats (allowing to load with weights_only=True), and saves the new snapshot data.

Warning: This overwrites your existing snapshot. If you want to ensure that no data is lost, copy your snapshot before calling fix_snapshot_metadata.

Parameters:

Name	Type	Description	Default
path	str | Path	The path of the snapshot to fix.	required

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

def fix_snapshot_metadata(path: str | Path) -> None:
    """Replace numpy floats in snapshot metrics.

    Only call this method with snapshots that you trust, as torch.load(...) is called
    with `weights_only=False`. For more information, see:
        https://pytorch.org/docs/stable/generated/torch.load.html

    DeepLabCut PyTorch snapshots trained with older releases may have `numpy` floats in
    the stored metrics. This method opens the snapshots (with `weights_only=False`),
    replaces the numpy floats with python floats (allowing to load with
    `weights_only=True`), and saves the new snapshot data.

    Warning: This overwrites your existing snapshot. If you want to ensure that no data
    is lost, copy your snapshot before calling `fix_snapshot_metadata`.

    Args:
        path: The path of the snapshot to fix.
    """
    snapshot = torch.load(path, map_location="cpu", weights_only=False)
    metrics = snapshot.get("metadata", {}).get("metrics")
    if metrics is not None:
        snapshot["metadata"]["metrics"] = {k: float(v) for k, v in metrics.items()}

    torch.save(snapshot, path)

get_load_weights_only

get_load_weights_only() -> bool

Gets the default value to use when loading snapshots with torch.load(...).

Returns:

Type	Description
bool	The default weights_only value when loading snapshots using torch.load(...).

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

def get_load_weights_only() -> bool:
    """Gets the default value to use when loading snapshots with `torch.load(...)`.

    Returns:
        The default `weights_only` value when loading snapshots using `torch.load(...)`.
    """
    global _load_weights_only
    return _load_weights_only

set_load_weights_only

set_load_weights_only(value: bool) -> None

Sets the default value to use when loading snapshots with torch.load(...).

Parameters:

Name	Type	Description	Default
value	bool	The default weights_only value to use when loading snapshots using torch.load(...).	required

Source code in deeplabcut/pose_estimation_pytorch/runners/base.py

def set_load_weights_only(value: bool) -> None:
    """Sets the default value to use when loading snapshots with `torch.load(...)`.

    Args:
        value: The default `weights_only` value to use when loading snapshots using
            `torch.load(...)`.
    """
    global _load_weights_only
    _load_weights_only = value