Skip to content

deeplabcut.pose_estimation_pytorch.models.heads.base

Classes:

Name Description
BaseHead

A head for pose estimation models.
WeightConversionMixin

A mixin for heads that can re-order and/or filter the output channels.

BaseHead

Bases: ABC, Module

A head for pose estimation models.

Attributes:

Name Type Description
stride

The stride for the head (or neck + head pair), where positive values indicate an increase in resolution while negative values a decrease. Assuming that H and W are divisible by stride, this is the value such that if a backbone outputs an encoding of shape (C, H, W), the head will output heatmaps of shape: (C, H * stride, W * stride) if stride > 0 (C, -H/stride, -W/stride) if stride < 0
predictor

an object to generate predictions from the head outputs
target_generator

a target generator which must output a target for each output key of this module (i.e. if forward returns a "heatmap" tensor and an "offset" tensor, then targets must be generated for both)
criterion

either a single criterion (e.g. if this head only outputs heatmaps) or a dictionary mapping the outputs of this head to the criterion to use (e.g. a "heatmap" criterion and an "offset" criterion for DEKR).
aggregator

if the criterion is a dictionary, cannot be none. used to combine the individual losses from this head into one "total_loss"

Methods:

Name Description
forward

Given the feature maps for an image ()
get_loss

Computes the loss for this head.
Source code in deeplabcut/pose_estimation_pytorch/models/heads/base.py 
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 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
class BaseHead(ABC, nn.Module):
    """A head for pose estimation models.

    Attributes:
        stride: The stride for the head (or neck + head pair), where positive values
            indicate an increase in resolution while negative values a decrease.
            Assuming that H and W are divisible by `stride`, this is the value such
            that if a backbone outputs an encoding of shape (C, H, W), the head will
            output heatmaps of shape:
                (C, H * stride, W * stride)       if stride > 0
                (C, -H/stride, -W/stride)         if stride < 0
        predictor: an object to generate predictions from the head outputs
        target_generator: a target generator which must output a target for each
            output key of this module (i.e. if forward returns a "heatmap" tensor and
            an "offset" tensor, then targets must be generated for both)
        criterion: either a single criterion (e.g. if this head only outputs heatmaps)
            or a dictionary mapping the outputs of this head to the criterion to use
            (e.g. a "heatmap" criterion and an "offset" criterion for DEKR).
        aggregator: if the criterion is a dictionary, cannot be none. used to combine
            the individual losses from this head into one "total_loss"
    """

    def __init__(
        self,
        stride: int | float,
        predictor: BasePredictor,
        target_generator: BaseGenerator,
        criterion: dict[str, BaseCriterion] | BaseCriterion,
        aggregator: BaseLossAggregator | None = None,
        weight_init: str | dict | BaseWeightInitializer | None = None,
    ) -> None:
        super().__init__()
        if stride == 0:
            raise ValueError(f"Stride must not be 0. Found {stride}.")

        self.stride = stride
        self.predictor = predictor
        self.target_generator = target_generator
        self.criterion = criterion
        self.aggregator = aggregator

        self.weight_init: BaseWeightInitializer | None = None
        if isinstance(weight_init, BaseWeightInitializer):
            self.weight_init = weight_init
        elif isinstance(weight_init, (str, dict)):
            self.weight_init = WEIGHT_INIT.build(weight_init)
        elif weight_init is not None:
            raise ValueError(f"Could not parse ``weight_init`` parameter: {weight_init}.")

        if isinstance(criterion, dict):
            if aggregator is None:
                raise ValueError("When multiple criterions are defined, a loss aggregator must also be given")
        else:
            if aggregator is not None:
                raise ValueError("Cannot use a loss aggregator with a single criterion")

    @abstractmethod
    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        """Given the feature maps for an image ()

        Args:
            x: the feature maps, of shape (b, c, h, w)

        Returns:
            the head outputs (e.g. "heatmap", "locref")
        """
        pass

    def get_loss(
        self,
        outputs: dict[str, torch.Tensor],
        targets: dict[str, dict[str, torch.Tensor]],
    ) -> dict[str, torch.Tensor]:
        """Computes the loss for this head.

        Args:
            outputs: the outputs of this head
            targets: the targets for this head

        Returns:
            A dictionary containing minimally "total_loss" key mapping to the total
            loss for this head (from which backwards() should be called). Can contain
            other keys containing losses that can be logged for informational purposes.
        """
        if self.aggregator is None:
            assert len(outputs) == len(targets) == 1
            key = [k for k in outputs.keys()][0]
            return {"total_loss": self.criterion(outputs[key], **targets[key])}

        losses = {name: criterion(outputs[name], **targets[name]) for name, criterion in self.criterion.items()}
        losses["total_loss"] = self.aggregator(losses)
        return losses

    def _init_weights(self) -> None:
        """Should be called once all modules for the class are created."""
        if self.weight_init is not None:
            self.weight_init.init_weights(self)

forward abstractmethod 

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

Given the feature maps for an image ()

Parameters:

Name Type Description Default

x

Tensor

the feature maps, of shape (b, c, h, w)

 required

Returns:

Type Description
dict[str, Tensor]

the head outputs (e.g. "heatmap", "locref")
Source code in deeplabcut/pose_estimation_pytorch/models/heads/base.py 
89
90
91
92
93
94
95
96
97
98
99
@abstractmethod
def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Given the feature maps for an image ()

    Args:
        x: the feature maps, of shape (b, c, h, w)

    Returns:
        the head outputs (e.g. "heatmap", "locref")
    """
    pass

get_loss 

get_loss(outputs: dict[str, Tensor], targets: dict[str, dict[str, Tensor]]) -> dict[str, torch.Tensor]

Computes the loss for this head.

Parameters:

Name Type Description Default

outputs

dict[str, Tensor]

the outputs of this head

 required

targets

dict[str, dict[str, Tensor]]

the targets for this head

 required

Returns:

Type Description
dict[str, Tensor]

A dictionary containing minimally "total_loss" key mapping to the total loss for this head (from which backwards() should be called). Can contain other keys containing losses that can be logged for informational purposes.
Source code in deeplabcut/pose_estimation_pytorch/models/heads/base.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def get_loss(
    self,
    outputs: dict[str, torch.Tensor],
    targets: dict[str, dict[str, torch.Tensor]],
) -> dict[str, torch.Tensor]:
    """Computes the loss for this head.

    Args:
        outputs: the outputs of this head
        targets: the targets for this head

    Returns:
        A dictionary containing minimally "total_loss" key mapping to the total
        loss for this head (from which backwards() should be called). Can contain
        other keys containing losses that can be logged for informational purposes.
    """
    if self.aggregator is None:
        assert len(outputs) == len(targets) == 1
        key = [k for k in outputs.keys()][0]
        return {"total_loss": self.criterion(outputs[key], **targets[key])}

    losses = {name: criterion(outputs[name], **targets[name]) for name, criterion in self.criterion.items()}
    losses["total_loss"] = self.aggregator(losses)
    return losses

WeightConversionMixin

Bases: ABC

A mixin for heads that can re-order and/or filter the output channels.

This mixin is useful to convert SuperAnimal model weights such that they can be used in downstream projects (either existing or new), where only a subset of keypoints are available (and where they might be re-ordered).

Methods:

Name Description
convert_weights

Converts pre-trained weights to be fine-tuned on another dataset.
Source code in deeplabcut/pose_estimation_pytorch/models/heads/base.py 
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
class WeightConversionMixin(ABC):
    """A mixin for heads that can re-order and/or filter the output channels.

    This mixin is useful to convert SuperAnimal model weights such that they can be used
    in downstream projects (either existing or new), where only a subset of keypoints
    are available (and where they might be re-ordered).
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

    @staticmethod
    @abstractmethod
    def convert_weights(
        state_dict: dict[str, torch.Tensor],
        module_prefix: str,
        conversion: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        """Converts pre-trained weights to be fine-tuned on another dataset.

        Args:
            state_dict: the state dict for the pre-trained model
            module_prefix: the prefix for weights in this head (e.g., 'heads.bodypart.')
            conversion: the mapping of old indices to new indices

        Examples:
            A SuperAnimal model was trained on the keypoints ["ear_left", "ear_right",
            "eye_left", "eye_right", "nose"]. A down-stream project has the bodyparts
            labeled ["nose", "eye_left", "eye_right"]. The SuperAnimal weights can be
            converted (to be used with the downstream project) with the following code:

                ``
                state_dict = torch.load(
                    snapshot_path, map_location=torch.device('cpu')
                )["model"]
                state_dict = HeadClass.convert_weights(
                    state_dict,
                    "heads.bodypart",
                    [4, 2, 3]
                )
                ``
        """
        pass

convert_weights abstractmethod staticmethod 

convert_weights(state_dict: dict[str, Tensor], module_prefix: str, conversion: Tensor) -> dict[str, torch.Tensor]

Converts pre-trained weights to be fine-tuned on another dataset.

Parameters:

Name Type Description Default

state_dict

dict[str, Tensor]

the state dict for the pre-trained model

 required

module_prefix

str

the prefix for weights in this head (e.g., 'heads.bodypart.')

 required

conversion

Tensor

the mapping of old indices to new indices

 required

Examples:

A SuperAnimal model was trained on the keypoints ["ear_left", "ear_right", "eye_left", "eye_right", "nose"]. A down-stream project has the bodyparts labeled ["nose", "eye_left", "eye_right"]. The SuperAnimal weights can be converted (to be used with the downstream project) with the following code:

``
state_dict = torch.load(
    snapshot_path, map_location=torch.device('cpu')
)["model"]
state_dict = HeadClass.convert_weights(
    state_dict,
    "heads.bodypart",
    [4, 2, 3]
)
``
Source code in deeplabcut/pose_estimation_pytorch/models/heads/base.py 
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
@staticmethod
@abstractmethod
def convert_weights(
    state_dict: dict[str, torch.Tensor],
    module_prefix: str,
    conversion: torch.Tensor,
) -> dict[str, torch.Tensor]:
    """Converts pre-trained weights to be fine-tuned on another dataset.

    Args:
        state_dict: the state dict for the pre-trained model
        module_prefix: the prefix for weights in this head (e.g., 'heads.bodypart.')
        conversion: the mapping of old indices to new indices

    Examples:
        A SuperAnimal model was trained on the keypoints ["ear_left", "ear_right",
        "eye_left", "eye_right", "nose"]. A down-stream project has the bodyparts
        labeled ["nose", "eye_left", "eye_right"]. The SuperAnimal weights can be
        converted (to be used with the downstream project) with the following code:

            ``
            state_dict = torch.load(
                snapshot_path, map_location=torch.device('cpu')
            )["model"]
            state_dict = HeadClass.convert_weights(
                state_dict,
                "heads.bodypart",
                [4, 2, 3]
            )
            ``
    """
    pass