deeplabcut.pose_estimation_pytorch.models.model

Classes:

Name	Description
PoseModel	A pose estimation model.

Functions:

Name	Description
filter_state_dict	Filters keys in the state dict for a module to only keep a given prefix. Removes

PoseModel

Bases: Module

A pose estimation model.

A pose estimation model is composed of a backbone, optionally a neck, and an arbitrary number of heads. Outputs are computed as follows:

Methods:

Name	Description
__init__	Args:
build	Args:
forward	Forward pass of the PoseModel.
get_predictions	Abstract method for the forward pass of the Predictor.
get_stride	Args:
get_target	Summary:

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

class PoseModel(nn.Module):
    """A pose estimation model.

    A pose estimation model is composed of a backbone, optionally a neck, and an
    arbitrary number of heads. Outputs are computed as follows:
    """

    def __init__(
        self,
        cfg: dict,
        backbone: BaseBackbone,
        heads: dict[str, BaseHead],
        neck: BaseNeck | None = None,
    ) -> None:
        """
        Args:
            cfg: configuration dictionary for the model.
            backbone: backbone network architecture.
            heads: the heads for the model
            neck: neck network architecture (default is None). Defaults to None.
        """
        super().__init__()
        self.cfg = cfg
        self.backbone = backbone
        self.heads = nn.ModuleDict(heads)
        self.neck = neck
        self.output_features = False

        self._strides = {name: _model_stride(self.backbone.stride, head.stride) for name, head in heads.items()}

    def forward(self, x: torch.Tensor, **backbone_kwargs) -> dict[str, dict[str, torch.Tensor]]:
        """Forward pass of the PoseModel.

        Args:
            x: input images

        Returns:
            Outputs of head groups
        """
        if x.dim() == 3:
            x = x[None, :]
        features = self.backbone(x, **backbone_kwargs)
        if self.neck:
            features = self.neck(features)

        outputs = {}
        if self.output_features:
            outputs["backbone"] = dict(features=features)

        for head_name, head in self.heads.items():
            outputs[head_name] = head(features)
        return outputs

    def get_loss(
        self,
        outputs: dict[str, dict[str, torch.Tensor]],
        targets: dict[str, dict[str, torch.Tensor]],
    ) -> dict[str, torch.Tensor]:
        total_losses = []
        losses: dict[str, torch.Tensor] = {}
        for name, head in self.heads.items():
            head_losses = head.get_loss(outputs[name], targets[name])
            total_losses.append(head_losses["total_loss"])
            for k, v in head_losses.items():
                losses[f"{name}_{k}"] = v

        # TODO: Different aggregation for multi-head loss?
        losses["total_loss"] = torch.mean(torch.stack(total_losses))
        return losses

    def get_target(
        self,
        outputs: dict[str, dict[str, torch.Tensor]],
        labels: dict,
    ) -> dict[str, dict]:
        """Summary:
        Get targets for model training.

        Args:
            outputs: output of each head group
            labels: dictionary of labels

        Returns:
            targets: dict of the targets for each model head group
        """
        return {
            name: head.target_generator(self._strides[name], outputs[name], labels) for name, head in self.heads.items()
        }

    def get_predictions(self, outputs: dict[str, dict[str, torch.Tensor]]) -> dict:
        """Abstract method for the forward pass of the Predictor.

        Args:
            outputs: outputs of the model heads

        Returns:
            A dictionary containing the predictions of each head group
        """
        predictions = {name: head.predictor(self._strides[name], outputs[name]) for name, head in self.heads.items()}
        if self.output_features:
            predictions["backbone"] = outputs["backbone"]

        return predictions

    def get_stride(self, head: str) -> int:
        """
        Args:
            head: The head for which to get the total stride.

        Returns:
            The total stride for the outputs of the head.

        Raises:
            ValueError: If there is no such head.
        """
        return self._strides[head]

    @staticmethod
    def build(
        cfg: dict,
        weight_init: None | WeightInitialization = None,
        pretrained_backbone: bool = False,
    ) -> PoseModel:
        """
        Args:
            cfg: The configuration of the model to build.
            weight_init: How model weights should be initialized. If None, ImageNet
                pre-trained backbone weights are loaded from Timm.
            pretrained_backbone: Whether to load an ImageNet-pretrained weights for
                the backbone. This should only be set to True when building a model
                which will be trained on a transfer learning task.

        Returns:
            the built pose model
        """
        cfg["backbone"]["pretrained"] = pretrained_backbone
        backbone = BACKBONES.build(dict(cfg["backbone"]))

        neck = None
        if cfg.get("neck"):
            neck = NECKS.build(dict(cfg["neck"]))

        heads = {}
        for name, head_cfg in cfg["heads"].items():
            head_cfg = copy.deepcopy(head_cfg)
            if "type" in head_cfg["criterion"]:
                head_cfg["criterion"] = CRITERIONS.build(head_cfg["criterion"])
            else:
                weights = {}
                criterions = {}
                for loss_name, criterion_cfg in head_cfg["criterion"].items():
                    weights[loss_name] = criterion_cfg.get("weight", 1.0)
                    criterion_cfg = {k: v for k, v in criterion_cfg.items() if k != "weight"}
                    criterions[loss_name] = CRITERIONS.build(criterion_cfg)

                aggregator_cfg = {"type": "WeightedLossAggregator", "weights": weights}
                head_cfg["aggregator"] = LOSS_AGGREGATORS.build(aggregator_cfg)
                head_cfg["criterion"] = criterions

            head_cfg["target_generator"] = TARGET_GENERATORS.build(head_cfg["target_generator"])
            head_cfg["predictor"] = PREDICTORS.build(head_cfg["predictor"])
            heads[name] = HEADS.build(head_cfg)

        model = PoseModel(cfg=cfg, backbone=backbone, neck=neck, heads=heads)

        if weight_init is not None:
            logging.info(f"Loading pretrained model weights: {weight_init}")
            logging.info(f"The pose model is loading from {weight_init.snapshot_path}")
            snapshot = torch.load(weight_init.snapshot_path, map_location="cpu")
            state_dict = snapshot["model"]

            # load backbone state dict
            model.backbone.load_state_dict(filter_state_dict(state_dict, "backbone"))

            # if there's a neck, load state dict
            if model.neck is not None:
                model.neck.load_state_dict(filter_state_dict(state_dict, "neck"))

            # load head state dicts
            if weight_init.with_decoder:
                all_head_state_dicts = filter_state_dict(state_dict, "heads")
                conversion_tensor = torch.from_numpy(weight_init.conversion_array)
                for name, head in model.heads.items():
                    head_state_dict = filter_state_dict(all_head_state_dicts, name)

                    # requires WeightConversionMixin
                    if not weight_init.memory_replay:
                        head_state_dict = head.convert_weights(
                            state_dict=head_state_dict,
                            module_prefix="",
                            conversion=conversion_tensor,
                        )

                    head.load_state_dict(head_state_dict)

        return model

__init__

__init__(cfg: dict, backbone: BaseBackbone, heads: dict[str, BaseHead], neck: BaseNeck | None = None) -> None

Parameters:

Name	Type	Description	Default
cfg	dict	configuration dictionary for the model.	required
backbone	BaseBackbone	backbone network architecture.	required
heads	dict[str, BaseHead]	the heads for the model	required
neck	BaseNeck | None	neck network architecture (default is None). Defaults to None.	None

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def __init__(
    self,
    cfg: dict,
    backbone: BaseBackbone,
    heads: dict[str, BaseHead],
    neck: BaseNeck | None = None,
) -> None:
    """
    Args:
        cfg: configuration dictionary for the model.
        backbone: backbone network architecture.
        heads: the heads for the model
        neck: neck network architecture (default is None). Defaults to None.
    """
    super().__init__()
    self.cfg = cfg
    self.backbone = backbone
    self.heads = nn.ModuleDict(heads)
    self.neck = neck
    self.output_features = False

    self._strides = {name: _model_stride(self.backbone.stride, head.stride) for name, head in heads.items()}

build staticmethod

build(cfg: dict, weight_init: None | WeightInitialization = None, pretrained_backbone: bool = False) -> PoseModel

Parameters:

Name	Type	Description	Default
cfg	dict	The configuration of the model to build.	required
weight_init	None | WeightInitialization	How model weights should be initialized. If None, ImageNet pre-trained backbone weights are loaded from Timm.	None
pretrained_backbone	bool	Whether to load an ImageNet-pretrained weights for the backbone. This should only be set to True when building a model which will be trained on a transfer learning task.	False

Returns:

Type	Description
PoseModel	the built pose model

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

@staticmethod
def build(
    cfg: dict,
    weight_init: None | WeightInitialization = None,
    pretrained_backbone: bool = False,
) -> PoseModel:
    """
    Args:
        cfg: The configuration of the model to build.
        weight_init: How model weights should be initialized. If None, ImageNet
            pre-trained backbone weights are loaded from Timm.
        pretrained_backbone: Whether to load an ImageNet-pretrained weights for
            the backbone. This should only be set to True when building a model
            which will be trained on a transfer learning task.

    Returns:
        the built pose model
    """
    cfg["backbone"]["pretrained"] = pretrained_backbone
    backbone = BACKBONES.build(dict(cfg["backbone"]))

    neck = None
    if cfg.get("neck"):
        neck = NECKS.build(dict(cfg["neck"]))

    heads = {}
    for name, head_cfg in cfg["heads"].items():
        head_cfg = copy.deepcopy(head_cfg)
        if "type" in head_cfg["criterion"]:
            head_cfg["criterion"] = CRITERIONS.build(head_cfg["criterion"])
        else:
            weights = {}
            criterions = {}
            for loss_name, criterion_cfg in head_cfg["criterion"].items():
                weights[loss_name] = criterion_cfg.get("weight", 1.0)
                criterion_cfg = {k: v for k, v in criterion_cfg.items() if k != "weight"}
                criterions[loss_name] = CRITERIONS.build(criterion_cfg)

            aggregator_cfg = {"type": "WeightedLossAggregator", "weights": weights}
            head_cfg["aggregator"] = LOSS_AGGREGATORS.build(aggregator_cfg)
            head_cfg["criterion"] = criterions

        head_cfg["target_generator"] = TARGET_GENERATORS.build(head_cfg["target_generator"])
        head_cfg["predictor"] = PREDICTORS.build(head_cfg["predictor"])
        heads[name] = HEADS.build(head_cfg)

    model = PoseModel(cfg=cfg, backbone=backbone, neck=neck, heads=heads)

    if weight_init is not None:
        logging.info(f"Loading pretrained model weights: {weight_init}")
        logging.info(f"The pose model is loading from {weight_init.snapshot_path}")
        snapshot = torch.load(weight_init.snapshot_path, map_location="cpu")
        state_dict = snapshot["model"]

        # load backbone state dict
        model.backbone.load_state_dict(filter_state_dict(state_dict, "backbone"))

        # if there's a neck, load state dict
        if model.neck is not None:
            model.neck.load_state_dict(filter_state_dict(state_dict, "neck"))

        # load head state dicts
        if weight_init.with_decoder:
            all_head_state_dicts = filter_state_dict(state_dict, "heads")
            conversion_tensor = torch.from_numpy(weight_init.conversion_array)
            for name, head in model.heads.items():
                head_state_dict = filter_state_dict(all_head_state_dicts, name)

                # requires WeightConversionMixin
                if not weight_init.memory_replay:
                    head_state_dict = head.convert_weights(
                        state_dict=head_state_dict,
                        module_prefix="",
                        conversion=conversion_tensor,
                    )

                head.load_state_dict(head_state_dict)

    return model

forward

forward(x: Tensor, **backbone_kwargs) -> dict[str, dict[str, torch.Tensor]]

Forward pass of the PoseModel.

Parameters:

Name	Type	Description	Default
x	Tensor	input images	required

Returns:

Type	Description
dict[str, dict[str, Tensor]]	Outputs of head groups

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def forward(self, x: torch.Tensor, **backbone_kwargs) -> dict[str, dict[str, torch.Tensor]]:
    """Forward pass of the PoseModel.

    Args:
        x: input images

    Returns:
        Outputs of head groups
    """
    if x.dim() == 3:
        x = x[None, :]
    features = self.backbone(x, **backbone_kwargs)
    if self.neck:
        features = self.neck(features)

    outputs = {}
    if self.output_features:
        outputs["backbone"] = dict(features=features)

    for head_name, head in self.heads.items():
        outputs[head_name] = head(features)
    return outputs

get_predictions

get_predictions(outputs: dict[str, dict[str, Tensor]]) -> dict

Abstract method for the forward pass of the Predictor.

Parameters:

Name	Type	Description	Default
outputs	dict[str, dict[str, Tensor]]	outputs of the model heads	required

Returns:

Type	Description
dict	A dictionary containing the predictions of each head group

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def get_predictions(self, outputs: dict[str, dict[str, torch.Tensor]]) -> dict:
    """Abstract method for the forward pass of the Predictor.

    Args:
        outputs: outputs of the model heads

    Returns:
        A dictionary containing the predictions of each head group
    """
    predictions = {name: head.predictor(self._strides[name], outputs[name]) for name, head in self.heads.items()}
    if self.output_features:
        predictions["backbone"] = outputs["backbone"]

    return predictions

get_stride

get_stride(head: str) -> int

Parameters:

Name	Type	Description	Default
head	str	The head for which to get the total stride.	required

Returns:

Type	Description
int	The total stride for the outputs of the head.

Raises:

Type	Description
ValueError	If there is no such head.

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def get_stride(self, head: str) -> int:
    """
    Args:
        head: The head for which to get the total stride.

    Returns:
        The total stride for the outputs of the head.

    Raises:
        ValueError: If there is no such head.
    """
    return self._strides[head]

get_target

get_target(outputs: dict[str, dict[str, Tensor]], labels: dict) -> dict[str, dict]

Summary: Get targets for model training.

Parameters:

Name	Type	Description	Default
outputs	dict[str, dict[str, Tensor]]	output of each head group	required
labels	dict	dictionary of labels	required

Returns:

Name	Type	Description
targets	dict[str, dict]	dict of the targets for each model head group

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def get_target(
    self,
    outputs: dict[str, dict[str, torch.Tensor]],
    labels: dict,
) -> dict[str, dict]:
    """Summary:
    Get targets for model training.

    Args:
        outputs: output of each head group
        labels: dictionary of labels

    Returns:
        targets: dict of the targets for each model head group
    """
    return {
        name: head.target_generator(self._strides[name], outputs[name], labels) for name, head in self.heads.items()
    }

filter_state_dict

filter_state_dict(state_dict: dict, module: str) -> dict[str, torch.Tensor]

Filters keys in the state dict for a module to only keep a given prefix. Removes the module from the keys (e.g. for module="backbone", "backbone.stage1.weight" will be converted to "stage1.weight" so the state dict can be loaded into the backbone directly).

Parameters:

Name	Type	Description	Default
state_dict	dict	the state dict	required
module	str	the module to keep, e.g. "backbone"	required

Returns:

Type	Description
dict[str, Tensor]	the filtered state dict, with the module removed from the keys

Examples:

state_dict = {"backbone.conv.weight": t1, "head.conv.weight": t2} filtered = filter_state_dict(state_dict, "backbone")

filtered =

model.backbone.load_state_dict(filtered)

Source code in deeplabcut/pose_estimation_pytorch/models/model.py

def filter_state_dict(state_dict: dict, module: str) -> dict[str, torch.Tensor]:
    """Filters keys in the state dict for a module to only keep a given prefix. Removes
    the module from the keys (e.g. for module="backbone", "backbone.stage1.weight" will
    be converted to "stage1.weight" so the state dict can be loaded into the backbone
    directly).

    Args:
        state_dict: the state dict
        module: the module to keep, e.g. "backbone"

    Returns:
        the filtered state dict, with the module removed from the keys

    Examples:
        state_dict = {"backbone.conv.weight": t1, "head.conv.weight": t2}
        filtered = filter_state_dict(state_dict, "backbone")
        # filtered = {"conv.weight": t1}
        model.backbone.load_state_dict(filtered)
    """
    return {
        ".".join(k.split(".")[1:]): v  # remove 'backbone.' from the keys
        for k, v in state_dict.items()
        if k.startswith(module)
    }