heads

sleap_nn.architectures.heads

Model head definitions for defining model output types.

Classes:

Name	Description
BBoxCenterHead	Center heatmap head of a bbox_centernet (CenterNet-style) detector.
BBoxOffsetHead	Sub-cell center offset head of a bbox_centernet.
BBoxSizeHead	Box size (width, height in output-grid cells) head of a bbox_centernet.
CenterOffsetHead	Head for predicting per-pixel offset vectors to instance centers.
CenteredInstanceConfmapsHead	Head for specifying centered instance confidence maps.
CentroidConfmapsHead	Head for specifying instance centroid confidence maps.
ClassMapsHead	Head for specifying class identity maps.
ClassVectorsHead	Head for specifying classification heads.
EmbeddingHead	Head for crop -> embedding-vector (re-ID) models.
GeM	Generalized-mean pooling: (mean(x.clamp(min=eps)^p))^(1/p) over HxW.
Head	Base class for model output heads.
InstanceCenterHead	Head for predicting instance center heatmaps.
L2Norm	L2-normalize along dim (so embeddings live on the unit hypersphere).
MultiInstanceConfmapsHead	Head for specifying multi-instance confidence maps.
MultiScaleHead	Base class for heads that read several backbone strides at once.
OffsetRefinementHead	Head for specifying offset refinement maps.
PartAffinityFieldsHead	Head for specifying multi-instance part affinity fields.
SSDHead	Multi-scale anchor head of a bbox_ssd (SSD-style) detector.
SegmentationHead	Head for predicting binary foreground segmentation masks.
SingleInstanceConfmapsHead	Head for specifying single instance confidence maps.

Functions:

Name	Description
default_embedding_pool	The pooling an embedding head uses when its config leaves pool unset.
embedding_encoder_is_frozen	Whether an embedding model's pooled feature comes from a frozen pretrained encoder.

BBoxCenterHead

Bases: Head

Center heatmap head of a bbox_centernet (CenterNet-style) detector.

One channel per box category; outputs logits (the sigmoid is applied by the loss and the inference forward) so the bias can take the focal-loss prior initialization.

Attributes:

Name	Type	Description
categories		Ordered class names, one channel each.
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class BBoxCenterHead(Head):
    """Center heatmap head of a ``bbox_centernet`` (CenterNet-style) detector.

    One channel per box category; outputs logits (the sigmoid is applied by the
    loss and the inference forward) so the bias can take the focal-loss prior
    initialization.

    Attributes:
        categories: Ordered class names, one channel each.
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        categories: Sequence[Text],
        output_stride: int = 4,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.categories = list(categories)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.categories)

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "centernet_focal"

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(categories, output_stride=4, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    categories: Sequence[Text],
    output_stride: int = 4,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.categories = list(categories)

BBoxOffsetHead

Bases: Head

Sub-cell center offset head of a bbox_centernet.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Attributes:

Name	Type	Description
channels	int	Return the number of channels in the tensor output by this head.
loss_function	str	Return the name of the loss function to use for this head.

Source code in sleap_nn/architectures/heads.py

class BBoxOffsetHead(Head):
    """Sub-cell center offset head of a ``bbox_centernet``."""

    def __init__(self, output_stride: int = 4, loss_weight: float = 1.0) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 2

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "masked_l1"

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=4, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(self, output_stride: int = 4, loss_weight: float = 1.0) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

BBoxSizeHead

Bases: Head

Box size (width, height in output-grid cells) head of a bbox_centernet.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Attributes:

Name	Type	Description
channels	int	Return the number of channels in the tensor output by this head.
loss_function	str	Return the name of the loss function to use for this head.

Source code in sleap_nn/architectures/heads.py

class BBoxSizeHead(Head):
    """Box size (width, height in output-grid cells) head of a ``bbox_centernet``."""

    def __init__(self, output_stride: int = 4, loss_weight: float = 0.1) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 2

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "masked_l1"

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=4, loss_weight=0.1)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(self, output_stride: int = 4, loss_weight: float = 0.1) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

CenterOffsetHead

Bases: Head

Head for predicting per-pixel offset vectors to instance centers.

Outputs a 2-channel map where each pixel's value is (dx, dy) pointing from the pixel to its instance's center. Only meaningful on foreground pixels.

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class CenterOffsetHead(Head):
    """Head for predicting per-pixel offset vectors to instance centers.

    Outputs a 2-channel map where each pixel's value is (dx, dy) pointing
    from the pixel to its instance's center. Only meaningful on foreground pixels.

    Attributes:
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        output_stride: int = 2,
        loss_weight: float = 0.1,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 2

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "smooth_l1"

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=2, loss_weight=0.1)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    output_stride: int = 2,
    loss_weight: float = 0.1,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

CenteredInstanceConfmapsHead

Bases: Head

Head for specifying centered instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
anchor_part		Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.
centroid_method		How the centroid is derived from the instance's points -- "center_of_mass", "bbox_center", "geometric_median" or "anchor"; None infers it from anchor_part. Data-pipeline metadata only (does not affect the head tensor); stored here so it rides the checkpoint and inference can reproduce the training geometry. See sleap_nn.data.instance_centroids.
centroid_fallback		Reduce method used when anchor_part is not visible. Passthrough metadata, as centroid_method.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class CenteredInstanceConfmapsHead(Head):
    """Head for specifying centered instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        anchor_part: Name of the part to use as an anchor node. If not specified, the
            bounding box centroid will be used.
        centroid_method: How the centroid is derived from the instance's points --
            ``"center_of_mass"``, ``"bbox_center"``, ``"geometric_median"`` or
            ``"anchor"``; ``None`` infers it from ``anchor_part``. Data-pipeline
            metadata only (does not affect the head tensor); stored here so it
            rides the checkpoint and inference can reproduce the training
            geometry. See ``sleap_nn.data.instance_centroids``.
        centroid_fallback: Reduce method used when ``anchor_part`` is not visible.
            Passthrough metadata, as ``centroid_method``.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        anchor_part: Optional[Text] = None,
        centroid_method: Optional[Text] = None,
        centroid_fallback: Optional[Text] = None,
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.anchor_part = anchor_part
        # Data-pipeline metadata only (does not affect the head tensor); see
        # `CentroidConfmapsHead` and `sleap_nn.data.instance_centroids`.
        self.centroid_method = centroid_method
        self.centroid_fallback = centroid_fallback
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "CenteredInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: Text name of the body parts (nodes) that the head will be
                configured to produce. The number of parts determines the number of
                channels in the output. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            anchor_part=config.anchor_part,
            centroid_method=config.get("centroid_method", None),
            centroid_fallback=config.get("centroid_fallback", None),
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, anchor_part=None, centroid_method=None, centroid_fallback=None, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    anchor_part: Optional[Text] = None,
    centroid_method: Optional[Text] = None,
    centroid_fallback: Optional[Text] = None,
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.anchor_part = anchor_part
    # Data-pipeline metadata only (does not affect the head tensor); see
    # `CentroidConfmapsHead` and `sleap_nn.data.instance_centroids`.
    self.centroid_method = centroid_method
    self.centroid_fallback = centroid_fallback
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		Text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
CenteredInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "CenteredInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: Text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of
            channels in the output. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        anchor_part=config.anchor_part,
        centroid_method=config.get("centroid_method", None),
        centroid_fallback=config.get("centroid_fallback", None),
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

CentroidConfmapsHead

Bases: Head

Head for specifying instance centroid confidence maps.

Attributes:

Name	Type	Description
anchor_part		Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.
centroid_source		Data-pipeline setting for which centroid the model is trained to predict ("user" / "computed" / None to infer). The head output is identical regardless; it is stored here only so the head can be built from the confmaps config, which co-locates it with anchor_part. See CentroidConfMapsConfig and resolve_centroid_source.
centroid_method		How the centroid is derived from the instance's points -- "center_of_mass", "bbox_center", "geometric_median" or "anchor"; None infers it from anchor_part. Data-pipeline metadata only (does not affect the head tensor); stored here so it rides the checkpoint and inference can reproduce the training geometry. See sleap_nn.data.instance_centroids.
centroid_fallback		Reduce method used when anchor_part is not visible. Passthrough metadata, as centroid_method.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
use_sigmoid_activation		If True, applies a sigmoid activation to the head's output so it is a calibrated (0, 1) probability rather than raw/unbounded regression output -- required when training with a focal-style loss (see CentroidConfMapsConfig.focal_loss_alpha), which needs Ŷ ∈ (0, 1) for its log(Ŷ)/log(1-Ŷ) terms, and so that peak_threshold-based inference (which compares raw confmap values against a fixed cutoff) stays meaningful. Default False (plain "identity" activation, i.e. no change from existing behavior) -- same opt-in pattern already used by SegmentationHead.
focal_loss_alpha		Training-loss metadata only (does not affect the head tensor); kept so the head is constructible from **head_config.confmaps. See CentroidConfMapsConfig.focal_loss_alpha -- actually consumed by CentroidLightningModule, not this head.
focal_loss_beta		Same as focal_loss_alpha -- passthrough only.
focal_loss_pos_threshold		Same as focal_loss_alpha -- passthrough only.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class CentroidConfmapsHead(Head):
    """Head for specifying instance centroid confidence maps.

    Attributes:
        anchor_part: Name of the part to use as an anchor node. If not specified, the
            bounding box centroid will be used.
        centroid_source: Data-pipeline setting for which centroid the model is
            trained to predict (``"user"`` / ``"computed"`` / ``None`` to infer).
            The head output is identical regardless; it is stored here only so
            the head can be built from the confmaps config, which co-locates it
            with ``anchor_part``. See ``CentroidConfMapsConfig`` and
            ``resolve_centroid_source``.
        centroid_method: How the centroid is derived from the instance's points --
            ``"center_of_mass"``, ``"bbox_center"``, ``"geometric_median"`` or
            ``"anchor"``; ``None`` infers it from ``anchor_part``. Data-pipeline
            metadata only (does not affect the head tensor); stored here so it
            rides the checkpoint and inference can reproduce the training
            geometry. See ``sleap_nn.data.instance_centroids``.
        centroid_fallback: Reduce method used when ``anchor_part`` is not visible.
            Passthrough metadata, as ``centroid_method``.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        use_sigmoid_activation: If `True`, applies a sigmoid activation to the head's
            output so it is a calibrated `(0, 1)` probability rather than raw/unbounded
            regression output -- required when training with a focal-style loss (see
            `CentroidConfMapsConfig.focal_loss_alpha`), which needs `Ŷ ∈ (0, 1)` for its
            `log(Ŷ)`/`log(1-Ŷ)` terms, and so that `peak_threshold`-based inference
            (which compares raw confmap values against a fixed cutoff) stays meaningful.
            Default `False` (plain "identity" activation, i.e. no change from existing
            behavior) -- same opt-in pattern already used by `SegmentationHead`.
        focal_loss_alpha: Training-loss metadata only (does not affect the head
            tensor); kept so the head is constructible from
            ``**head_config.confmaps``. See ``CentroidConfMapsConfig.focal_loss_alpha``
            -- actually consumed by ``CentroidLightningModule``, not this head.
        focal_loss_beta: Same as ``focal_loss_alpha`` -- passthrough only.
        focal_loss_pos_threshold: Same as ``focal_loss_alpha`` -- passthrough only.
    """

    def __init__(
        self,
        anchor_part: Optional[Text] = None,
        centroid_source: Optional[Text] = None,
        centroid_method: Optional[Text] = None,
        centroid_fallback: Optional[Text] = None,
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        use_sigmoid_activation: bool = False,
        focal_loss_alpha: float = 0.0,
        focal_loss_beta: float = 4.0,
        focal_loss_pos_threshold: float = 0.5,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.anchor_part = anchor_part
        # Data-pipeline metadata only (does not affect the head tensor); kept so
        # the head is constructible from ``**head_config.confmaps``.
        self.centroid_source = centroid_source
        self.centroid_method = centroid_method
        self.centroid_fallback = centroid_fallback
        self.sigma = sigma
        self.use_sigmoid_activation = use_sigmoid_activation
        # Training-loss metadata only (consumed by `CentroidLightningModule`, not
        # this head) -- kept here purely so the head is constructible from
        # `**head_config.confmaps`, which co-locates them with `use_sigmoid_activation`.
        self.focal_loss_alpha = focal_loss_alpha
        self.focal_loss_beta = focal_loss_beta
        self.focal_loss_pos_threshold = focal_loss_pos_threshold

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "sigmoid" if self.use_sigmoid_activation else "identity"

    @classmethod
    def from_config(cls, config: DictConfig) -> "CentroidConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.

        Returns:
            The instantiated head with the specified configuration options.
        """
        return cls(
            anchor_part=config.anchor_part,
            centroid_source=config.get("centroid_source", None),
            centroid_method=config.get("centroid_method", None),
            centroid_fallback=config.get("centroid_fallback", None),
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
            use_sigmoid_activation=getattr(config, "use_sigmoid_activation", False),
            focal_loss_alpha=getattr(config, "focal_loss_alpha", 0.0),
            focal_loss_beta=getattr(config, "focal_loss_beta", 4.0),
            focal_loss_pos_threshold=getattr(config, "focal_loss_pos_threshold", 0.5),
        )

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

__init__(anchor_part=None, centroid_source=None, centroid_method=None, centroid_fallback=None, sigma=5.0, output_stride=1, loss_weight=1.0, use_sigmoid_activation=False, focal_loss_alpha=0.0, focal_loss_beta=4.0, focal_loss_pos_threshold=0.5)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    anchor_part: Optional[Text] = None,
    centroid_source: Optional[Text] = None,
    centroid_method: Optional[Text] = None,
    centroid_fallback: Optional[Text] = None,
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    use_sigmoid_activation: bool = False,
    focal_loss_alpha: float = 0.0,
    focal_loss_beta: float = 4.0,
    focal_loss_pos_threshold: float = 0.5,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.anchor_part = anchor_part
    # Data-pipeline metadata only (does not affect the head tensor); kept so
    # the head is constructible from ``**head_config.confmaps``.
    self.centroid_source = centroid_source
    self.centroid_method = centroid_method
    self.centroid_fallback = centroid_fallback
    self.sigma = sigma
    self.use_sigmoid_activation = use_sigmoid_activation
    # Training-loss metadata only (consumed by `CentroidLightningModule`, not
    # this head) -- kept here purely so the head is constructible from
    # `**head_config.confmaps`, which co-locates them with `use_sigmoid_activation`.
    self.focal_loss_alpha = focal_loss_alpha
    self.focal_loss_beta = focal_loss_beta
    self.focal_loss_pos_threshold = focal_loss_pos_threshold

from_config(config) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.

Returns:

Type	Description
CentroidConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(cls, config: DictConfig) -> "CentroidConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.

    Returns:
        The instantiated head with the specified configuration options.
    """
    return cls(
        anchor_part=config.anchor_part,
        centroid_source=config.get("centroid_source", None),
        centroid_method=config.get("centroid_method", None),
        centroid_fallback=config.get("centroid_fallback", None),
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
        use_sigmoid_activation=getattr(config, "use_sigmoid_activation", False),
        focal_loss_alpha=getattr(config, "focal_loss_alpha", 0.0),
        focal_loss_beta=getattr(config, "focal_loss_beta", 4.0),
        focal_loss_pos_threshold=getattr(config, "focal_loss_pos_threshold", 0.5),
    )

ClassMapsHead

Bases: Head

Head for specifying class identity maps.

Attributes:

Name	Type	Description
classes		List of string names of the classes.
sigma		Spread of the class maps around each node.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
class_output		How the classes map to sleap-io objects ("track" / "identity"). Carried through verbatim from the head config; unused by the architecture/loss. Accepted so the config can be splatted into the constructor. Appended AFTER the pre-existing arguments so a positional ClassMapsHead(classes, 5.0) still binds sigma.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class ClassMapsHead(Head):
    """Head for specifying class identity maps.

    Attributes:
        classes: List of string names of the classes.
        sigma: Spread of the class maps around each node.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        class_output: How the classes map to sleap-io objects (``"track"`` /
            ``"identity"``). Carried through verbatim from the head config; unused
            by the architecture/loss. Accepted so the config can be splatted into
            the constructor. Appended AFTER the pre-existing arguments so a
            positional ``ClassMapsHead(classes, 5.0)`` still binds ``sigma``.
    """

    def __init__(
        self,
        classes: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        class_output: str = "track",
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.classes = classes
        self.class_output = class_output
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.classes)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "sigmoid"

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        classes: Optional[List[Text]] = None,
    ) -> "ClassMapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            classes: List of string names of the classes that this head will predict.
                This must be set if the `classes` attribute of the configuration is not
                set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.classes is not None:
            classes = config.classes
        return cls(
            classes=classes,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

__init__(classes, sigma=5.0, output_stride=1, loss_weight=1.0, class_output='track')

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    classes: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    class_output: str = "track",
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.classes = classes
    self.class_output = class_output
    self.sigma = sigma

from_config(config, classes=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
classes		List of string names of the classes that this head will predict. This must be set if the classes attribute of the configuration is not set.

Returns:

Type	Description
ClassMapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    classes: Optional[List[Text]] = None,
) -> "ClassMapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        classes: List of string names of the classes that this head will predict.
            This must be set if the `classes` attribute of the configuration is not
            set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.classes is not None:
        classes = config.classes
    return cls(
        classes=classes,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

ClassVectorsHead

Bases: Head

Head for specifying classification heads.

Attributes:

Name	Type	Description
classes		List of string names of the classes.
num_fc_layers		Number of fully connected layers after flattening input features.
num_fc_units		Number of units (dimensions) in fully connected layers prior to classification output.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
class_output		How the classes map to sleap-io objects ("track" / "identity"). Carried through verbatim from the head config; unused by the architecture/loss. Accepted so the config can be splatted into the constructor. Appended AFTER the pre-existing arguments so a positional ClassVectorsHead(classes, 2) still binds num_fc_layers.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.
make_head	Make head output tensor from input feature tensor.

Source code in sleap_nn/architectures/heads.py

class ClassVectorsHead(Head):
    """Head for specifying classification heads.

    Attributes:
        classes: List of string names of the classes.
        num_fc_layers: Number of fully connected layers after flattening input features.
        num_fc_units: Number of units (dimensions) in fully connected layers prior to
            classification output.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        class_output: How the classes map to sleap-io objects (``"track"`` /
            ``"identity"``). Carried through verbatim from the head config; unused
            by the architecture/loss. Accepted so the config can be splatted into
            the constructor. Appended AFTER the pre-existing arguments so a
            positional ``ClassVectorsHead(classes, 2)`` still binds
            ``num_fc_layers``.
    """

    def __init__(
        self,
        classes: List[Text],
        num_fc_layers: int = 1,
        num_fc_units: int = 64,
        global_pool: bool = True,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        class_output: str = "track",
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.classes = classes
        self.class_output = class_output
        self.num_fc_layers = num_fc_layers
        self.num_fc_units = num_fc_units
        self.global_pool = global_pool

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.classes)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "softmax"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "categorical_crossentropy"

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        classes: Optional[List[Text]] = None,
    ) -> "ClassVectorsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            classes: List of string names of the classes that this head will predict.
                This must be set if the `classes` attribute of the configuration is not
                set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.classes is not None:
            classes = config.classes
        return cls(
            classes=classes,
            num_fc_layers=config.num_fc_layers,
            num_fc_units=config.num_fc_units,
            global_pool=config.global_pool,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make head output tensor from input feature tensor.

        Args:
            x_in: An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).

        Returns:
            A `nn.Sequential` with the correct shape for the head.
        """
        from collections import OrderedDict

        module_dict = OrderedDict()

        if self.global_pool:
            module_dict[f"pre_classification_global_pool"] = nn.AdaptiveMaxPool2d(1)

        module_dict[f"pre_classification_flatten"] = nn.Flatten(start_dim=1)

        for i in range(self.num_fc_layers):
            if i == 0:
                module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                    x_in, self.num_fc_units
                )
            else:
                module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                    self.num_fc_units, self.num_fc_units
                )
            module_dict[f"pre_classification{i}_relu"] = get_act_fn("relu")

        module_dict[f"ClassVectorsHead"] = nn.Linear(self.num_fc_units, self.channels)
        module_dict[f"softmax"] = get_act_fn("softmax")

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(classes, num_fc_layers=1, num_fc_units=64, global_pool=True, output_stride=1, loss_weight=1.0, class_output='track')

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    classes: List[Text],
    num_fc_layers: int = 1,
    num_fc_units: int = 64,
    global_pool: bool = True,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    class_output: str = "track",
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.classes = classes
    self.class_output = class_output
    self.num_fc_layers = num_fc_layers
    self.num_fc_units = num_fc_units
    self.global_pool = global_pool

from_config(config, classes=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
classes		List of string names of the classes that this head will predict. This must be set if the classes attribute of the configuration is not set.

Returns:

Type	Description
ClassVectorsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    classes: Optional[List[Text]] = None,
) -> "ClassVectorsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        classes: List of string names of the classes that this head will predict.
            This must be set if the `classes` attribute of the configuration is not
            set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.classes is not None:
        classes = config.classes
    return cls(
        classes=classes,
        num_fc_layers=config.num_fc_layers,
        num_fc_units=config.num_fc_units,
        global_pool=config.global_pool,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

make_head(x_in)

Make head output tensor from input feature tensor.

Parameters:

Name	Type	Description	Default
x_in	int	An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).	required

Returns:

Type	Description
Sequential	A nn.Sequential with the correct shape for the head.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make head output tensor from input feature tensor.

    Args:
        x_in: An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).

    Returns:
        A `nn.Sequential` with the correct shape for the head.
    """
    from collections import OrderedDict

    module_dict = OrderedDict()

    if self.global_pool:
        module_dict[f"pre_classification_global_pool"] = nn.AdaptiveMaxPool2d(1)

    module_dict[f"pre_classification_flatten"] = nn.Flatten(start_dim=1)

    for i in range(self.num_fc_layers):
        if i == 0:
            module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                x_in, self.num_fc_units
            )
        else:
            module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                self.num_fc_units, self.num_fc_units
            )
        module_dict[f"pre_classification{i}_relu"] = get_act_fn("relu")

    module_dict[f"ClassVectorsHead"] = nn.Linear(self.num_fc_units, self.channels)
    module_dict[f"softmax"] = get_act_fn("softmax")

    return nn.Sequential(module_dict)

EmbeddingHead

Bases: Head

Head for crop -> embedding-vector (re-ID) models.

Mirrors ClassVectorsHead (a pooled, non-spatial head): [pool] -> Flatten -> num_fc_layers x (Linear+ReLU) -> Linear(embedding_dim) -> [L2Norm]. The pooled feature comes from the backbone's middle_output (lone head, empty decoder).

Attributes:

Name	Type	Description
embedding_dim		Output embedding dimensionality.
num_fc_layers		Number of FC layers before the embedding output.
num_fc_units		Units in the pre-embedding FC layers.
pool		Pooling over the encoder feature map: gem | max | avg. None means "the default for this backbone"; Model resolves it with default_embedding_pool before building the head.
normalize		L2-normalize the output embedding.
output_stride		Should equal the backbone max_stride (so the decoder is empty and the head taps middle_output).
loss_weight		Weight of the loss term for this head.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a head-leaf configuration.
make_head	Make the head output module from the pooled-feature input channels.

Source code in sleap_nn/architectures/heads.py

class EmbeddingHead(Head):
    """Head for crop -> embedding-vector (re-ID) models.

    Mirrors ``ClassVectorsHead`` (a pooled, non-spatial head): ``[pool] -> Flatten ->
    num_fc_layers x (Linear+ReLU) -> Linear(embedding_dim) -> [L2Norm]``. The pooled
    feature comes from the backbone's ``middle_output`` (lone head, empty decoder).

    Attributes:
        embedding_dim: Output embedding dimensionality.
        num_fc_layers: Number of FC layers before the embedding output.
        num_fc_units: Units in the pre-embedding FC layers.
        pool: Pooling over the encoder feature map: ``gem`` | ``max`` | ``avg``.
            ``None`` means "the default for this backbone"; `Model` resolves it
            with ``default_embedding_pool`` before building the head.
        normalize: L2-normalize the output embedding.
        output_stride: Should equal the backbone max_stride (so the decoder is empty
            and the head taps ``middle_output``).
        loss_weight: Weight of the loss term for this head.
    """

    def __init__(
        self,
        embedding_dim: int = 128,
        num_fc_layers: int = 1,
        num_fc_units: int = 256,
        pool: str = "gem",
        normalize: bool = True,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.embedding_dim = embedding_dim
        self.num_fc_layers = num_fc_layers
        self.num_fc_units = num_fc_units
        self.pool = pool
        self.normalize = normalize

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return self.embedding_dim

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the loss-function name (informational).

        The contrastive loss is batch-level and is driven by the embedding
        LightningModule's ``objective``, not by this string.
        """
        return "supcon"

    @classmethod
    def from_config(cls, config: DictConfig) -> "EmbeddingHead":
        """Create this head from a head-leaf configuration."""
        return cls(
            embedding_dim=config.embedding_dim,
            num_fc_layers=config.num_fc_layers,
            num_fc_units=config.num_fc_units,
            pool=config.pool,
            normalize=config.normalize,
            output_stride=config.output_stride,
        )

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make the head output module from the pooled-feature input channels.

        Args:
            x_in: Number of channels of the encoder feature map (its channel dim).

        Returns:
            An ``nn.Sequential`` mapping ``[B, x_in, H, W] -> [B, embedding_dim]``.
        """
        module_dict = OrderedDict()
        if self.pool == "gem":
            module_dict["pre_embedding_pool"] = GeM()
        elif self.pool == "max":
            module_dict["pre_embedding_pool"] = nn.AdaptiveMaxPool2d(1)
        elif self.pool == "avg":
            module_dict["pre_embedding_pool"] = nn.AdaptiveAvgPool2d(1)
        else:
            message = f"Unknown pool '{self.pool}'; choose one of gem|max|avg."
            logger.error(message)
            raise ValueError(message)

        module_dict["pre_embedding_flatten"] = nn.Flatten(start_dim=1)

        d_in = x_in
        for i in range(self.num_fc_layers):
            module_dict[f"pre_embedding{i}_fc"] = nn.Linear(d_in, self.num_fc_units)
            module_dict[f"pre_embedding{i}_relu"] = get_act_fn("relu")
            d_in = self.num_fc_units

        module_dict["EmbeddingHead"] = nn.Linear(d_in, self.embedding_dim)
        if self.normalize:
            module_dict["l2norm"] = L2Norm(dim=1)

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the loss-function name (informational).

The contrastive loss is batch-level and is driven by the embedding LightningModule's objective, not by this string.

__init__(embedding_dim=128, num_fc_layers=1, num_fc_units=256, pool='gem', normalize=True, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    embedding_dim: int = 128,
    num_fc_layers: int = 1,
    num_fc_units: int = 256,
    pool: str = "gem",
    normalize: bool = True,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.embedding_dim = embedding_dim
    self.num_fc_layers = num_fc_layers
    self.num_fc_units = num_fc_units
    self.pool = pool
    self.normalize = normalize

from_config(config) classmethod

Create this head from a head-leaf configuration.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(cls, config: DictConfig) -> "EmbeddingHead":
    """Create this head from a head-leaf configuration."""
    return cls(
        embedding_dim=config.embedding_dim,
        num_fc_layers=config.num_fc_layers,
        num_fc_units=config.num_fc_units,
        pool=config.pool,
        normalize=config.normalize,
        output_stride=config.output_stride,
    )

make_head(x_in)

Make the head output module from the pooled-feature input channels.

Parameters:

Name	Type	Description	Default
x_in	int	Number of channels of the encoder feature map (its channel dim).	required

Returns:

Type	Description
Sequential	An nn.Sequential mapping [B, x_in, H, W] -> [B, embedding_dim].

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make the head output module from the pooled-feature input channels.

    Args:
        x_in: Number of channels of the encoder feature map (its channel dim).

    Returns:
        An ``nn.Sequential`` mapping ``[B, x_in, H, W] -> [B, embedding_dim]``.
    """
    module_dict = OrderedDict()
    if self.pool == "gem":
        module_dict["pre_embedding_pool"] = GeM()
    elif self.pool == "max":
        module_dict["pre_embedding_pool"] = nn.AdaptiveMaxPool2d(1)
    elif self.pool == "avg":
        module_dict["pre_embedding_pool"] = nn.AdaptiveAvgPool2d(1)
    else:
        message = f"Unknown pool '{self.pool}'; choose one of gem|max|avg."
        logger.error(message)
        raise ValueError(message)

    module_dict["pre_embedding_flatten"] = nn.Flatten(start_dim=1)

    d_in = x_in
    for i in range(self.num_fc_layers):
        module_dict[f"pre_embedding{i}_fc"] = nn.Linear(d_in, self.num_fc_units)
        module_dict[f"pre_embedding{i}_relu"] = get_act_fn("relu")
        d_in = self.num_fc_units

    module_dict["EmbeddingHead"] = nn.Linear(d_in, self.embedding_dim)
    if self.normalize:
        module_dict["l2norm"] = L2Norm(dim=1)

    return nn.Sequential(module_dict)

GeM

Bases: Module

Generalized-mean pooling: (mean(x.clamp(min=eps)^p))^(1/p) over HxW.

The exponent p is learnable (init 3.0). The clamp(min=eps) BEFORE the fractional power guards against NaNs (a fractional power of a negative/zero base). Returns a flattened [B, C] tensor.

The forward pass uses p clamped to [GEM_P_MIN, GEM_P_MAX]; the stored parameter is the raw exponent (same state-dict key and meaning as before the upper bound existed, so a checkpoint whose p lies in that range produces exactly the output it did). During training the clamp passes its gradient straight through (_ClampStraightThrough), so p cannot get stuck at a bound.

The eps floor also means a NEGATIVE activation contributes as if it were eps: GeM pools only the positive part of each channel. That is what the native backbones feed it (their bottleneck ends in a ReLU). A pretrained encoder ends in a LayerNorm instead (about half its activations are negative); when it is fine-tuned it adapts and GeM costs nothing, but a FROZEN one cannot, so the default pool for a frozen pretrained encoder is avg -- see default_embedding_pool.

Methods:

Name	Description
__init__	Initialize the pooling layer.
exponent	The exponent the forward pass uses: p clamped to the valid range.
forward	Pool [B, C, H, W] to [B, C] by the generalized mean over HxW.

Source code in sleap_nn/architectures/heads.py

class GeM(nn.Module):
    """Generalized-mean pooling: ``(mean(x.clamp(min=eps)^p))^(1/p)`` over HxW.

    The exponent ``p`` is learnable (init 3.0). The ``clamp(min=eps)`` BEFORE the
    fractional power guards against NaNs (a fractional power of a negative/zero base).
    Returns a flattened ``[B, C]`` tensor.

    The forward pass uses ``p`` clamped to ``[GEM_P_MIN, GEM_P_MAX]``; the stored
    parameter is the raw exponent (same state-dict key and meaning as before the
    upper bound existed, so a checkpoint whose ``p`` lies in that range produces
    exactly the output it did). During training the clamp passes its gradient
    straight through (``_ClampStraightThrough``), so ``p`` cannot get stuck at a
    bound.

    The ``eps`` floor also means a NEGATIVE activation contributes as if it were
    ``eps``: GeM pools only the positive part of each channel. That is what the
    native backbones feed it (their bottleneck ends in a ReLU). A pretrained encoder
    ends in a LayerNorm instead (about half its activations are negative); when it
    is fine-tuned it adapts and GeM costs nothing, but a FROZEN one cannot, so the
    default pool for a frozen pretrained encoder is ``avg`` -- see
    ``default_embedding_pool``.
    """

    def __init__(self, p: float = 3.0, eps: float = 1e-6, learnable: bool = True):
        """Initialize the pooling layer.

        Args:
            p: Initial generalized-mean exponent (``p=1`` averages, larger ``p``
                approaches max-pooling).
            eps: Small floor applied to the activation before the power to avoid a
                fractional power of a non-positive value.
            learnable: If ``True``, ``p`` is a trainable parameter; otherwise it is a
                fixed buffer.
        """
        super().__init__()
        if learnable:
            self.p = nn.Parameter(torch.tensor(float(p)))
        else:
            self.register_buffer("p", torch.tensor(float(p)))
        self.eps = eps

    def exponent(self) -> torch.Tensor:
        """The exponent the forward pass uses: ``p`` clamped to the valid range.

        Training (train mode, autograd on, learnable ``p``) routes the clamp through
        ``_ClampStraightThrough``; everywhere else it is a plain clamp with the same
        value, which is also what ONNX export traces.
        """
        if self.training and torch.is_grad_enabled() and self.p.requires_grad:
            return _ClampStraightThrough.apply(self.p, GEM_P_MIN, GEM_P_MAX)
        return self.p.clamp(GEM_P_MIN, GEM_P_MAX)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Pool ``[B, C, H, W]`` to ``[B, C]`` by the generalized mean over ``HxW``."""
        p = self.exponent()
        xp = x.clamp(min=self.eps).pow(p)
        return F.adaptive_avg_pool2d(xp, 1).pow(1.0 / p).flatten(1)

    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
        """Load as usual; warn if a stored exponent lies above the forward cap."""
        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
        stored = state_dict.get(prefix + "p")
        if stored is not None and float(stored) > GEM_P_MAX:
            logger.warning(
                f"GeM exponent p={float(stored):.4g} in this checkpoint is above the "
                f"cap {GEM_P_MAX}; pooling now uses p={GEM_P_MAX}, so its embeddings "
                "differ from the ones it was trained with."
            )

__init__(p=3.0, eps=1e-06, learnable=True)

Initialize the pooling layer.

Parameters:

Name	Type	Description	Default
p	float	Initial generalized-mean exponent (p=1 averages, larger p approaches max-pooling).	3.0
eps	float	Small floor applied to the activation before the power to avoid a fractional power of a non-positive value.	1e-06
learnable	bool	If True, p is a trainable parameter; otherwise it is a fixed buffer.	True

Source code in sleap_nn/architectures/heads.py

def __init__(self, p: float = 3.0, eps: float = 1e-6, learnable: bool = True):
    """Initialize the pooling layer.

    Args:
        p: Initial generalized-mean exponent (``p=1`` averages, larger ``p``
            approaches max-pooling).
        eps: Small floor applied to the activation before the power to avoid a
            fractional power of a non-positive value.
        learnable: If ``True``, ``p`` is a trainable parameter; otherwise it is a
            fixed buffer.
    """
    super().__init__()
    if learnable:
        self.p = nn.Parameter(torch.tensor(float(p)))
    else:
        self.register_buffer("p", torch.tensor(float(p)))
    self.eps = eps

exponent()

The exponent the forward pass uses: p clamped to the valid range.

Training (train mode, autograd on, learnable p) routes the clamp through _ClampStraightThrough; everywhere else it is a plain clamp with the same value, which is also what ONNX export traces.

Source code in sleap_nn/architectures/heads.py

def exponent(self) -> torch.Tensor:
    """The exponent the forward pass uses: ``p`` clamped to the valid range.

    Training (train mode, autograd on, learnable ``p``) routes the clamp through
    ``_ClampStraightThrough``; everywhere else it is a plain clamp with the same
    value, which is also what ONNX export traces.
    """
    if self.training and torch.is_grad_enabled() and self.p.requires_grad:
        return _ClampStraightThrough.apply(self.p, GEM_P_MIN, GEM_P_MAX)
    return self.p.clamp(GEM_P_MIN, GEM_P_MAX)

forward(x)

Pool [B, C, H, W] to [B, C] by the generalized mean over HxW.

Source code in sleap_nn/architectures/heads.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Pool ``[B, C, H, W]`` to ``[B, C]`` by the generalized mean over ``HxW``."""
    p = self.exponent()
    xp = x.clamp(min=self.eps).pow(p)
    return F.adaptive_avg_pool2d(xp, 1).pow(1.0 / p).flatten(1)

Head

Base class for model output heads.

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
make_head	Make head output tensor from input feature tensor.

Source code in sleap_nn/architectures/heads.py

class Head:
    """Base class for model output heads.

    Attributes:
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(self, output_stride: int = 1, loss_weight: float = 1.0) -> None:
        """Initialize the object with the specified attributes."""
        self.output_stride = output_stride
        self.loss_weight = loss_weight

    @property
    def name(self) -> str:
        """Name of the head."""
        return type(self).__name__

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        message = "Subclasses must implement this method."
        logger.error(message)
        raise NotImplementedError(message)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "mse"

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make head output tensor from input feature tensor.

        Args:
            x_in: An int input for the input channels.

        Returns:
            A `nn.Sequential` with the correct shape for the head.
        """
        module_dict = OrderedDict()
        module_dict[self.name] = nn.Sequential(
            nn.Conv2d(
                in_channels=x_in,
                out_channels=self.channels,
                kernel_size=1,
                stride=1,
                padding="same",
            ),
            get_act_fn(self.activation),
        )

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

name property

Name of the head.

__init__(output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(self, output_stride: int = 1, loss_weight: float = 1.0) -> None:
    """Initialize the object with the specified attributes."""
    self.output_stride = output_stride
    self.loss_weight = loss_weight

make_head(x_in)

Make head output tensor from input feature tensor.

Parameters:

Name	Type	Description	Default
x_in	int	An int input for the input channels.	required

Returns:

Type	Description
Sequential	A nn.Sequential with the correct shape for the head.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make head output tensor from input feature tensor.

    Args:
        x_in: An int input for the input channels.

    Returns:
        A `nn.Sequential` with the correct shape for the head.
    """
    module_dict = OrderedDict()
    module_dict[self.name] = nn.Sequential(
        nn.Conv2d(
            in_channels=x_in,
            out_channels=self.channels,
            kernel_size=1,
            stride=1,
            padding="same",
        ),
        get_act_fn(self.activation),
    )

    return nn.Sequential(module_dict)

InstanceCenterHead

Bases: Head

Head for predicting instance center heatmaps.

Outputs a single-channel Gaussian heatmap with peaks at each instance's mask centroid. Similar to CentroidConfmapsHead but for mask-derived centers.

Attributes:

Name	Type	Description
sigma		Standard deviation of the Gaussian in pixels.
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class InstanceCenterHead(Head):
    """Head for predicting instance center heatmaps.

    Outputs a single-channel Gaussian heatmap with peaks at each instance's
    mask centroid. Similar to CentroidConfmapsHead but for mask-derived centers.

    Attributes:
        sigma: Standard deviation of the Gaussian in pixels.
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        sigma: float = 4.0,
        output_stride: int = 2,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

channels property

Return the number of channels in the tensor output by this head.

__init__(sigma=4.0, output_stride=2, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    sigma: float = 4.0,
    output_stride: int = 2,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.sigma = sigma

L2Norm

Bases: Module

L2-normalize along dim (so embeddings live on the unit hypersphere).

Methods:

Name	Description
__init__	Initialize the layer.
forward	L2-normalize x along dim.

Source code in sleap_nn/architectures/heads.py

class L2Norm(nn.Module):
    """L2-normalize along ``dim`` (so embeddings live on the unit hypersphere)."""

    def __init__(self, dim: int = 1):
        """Initialize the layer.

        Args:
            dim: Dimension along which to L2-normalize.
        """
        super().__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """L2-normalize ``x`` along ``dim``."""
        return F.normalize(x, dim=self.dim)

__init__(dim=1)

Initialize the layer.

Parameters:

Name	Type	Description	Default
dim	int	Dimension along which to L2-normalize.	1

Source code in sleap_nn/architectures/heads.py

def __init__(self, dim: int = 1):
    """Initialize the layer.

    Args:
        dim: Dimension along which to L2-normalize.
    """
    super().__init__()
    self.dim = dim

forward(x)

L2-normalize x along dim.

Source code in sleap_nn/architectures/heads.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """L2-normalize ``x`` along ``dim``."""
    return F.normalize(x, dim=self.dim)

MultiInstanceConfmapsHead

Bases: Head

Head for specifying multi-instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class MultiInstanceConfmapsHead(Head):
    """Head for specifying multi-instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "MultiInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: Text name of the body parts (nodes) that the head will be
                configured to produce. The number of parts determines the number of
                channels in the output. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		Text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
MultiInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "MultiInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: Text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of
            channels in the output. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

MultiScaleHead

Bases: Head

Base class for heads that read several backbone strides at once.

Model binds a MultiScaleHead to the feature map at every stride in output_strides (decoder outputs, plus the bottleneck for the deepest one) and calls its layer with the list of those maps, finest first. output_stride is the finest of them.

Attributes:

Name	Type	Description
output_strides		Feature strides to read, finest first.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
make_head	Make the head from the input channels of each stride, finest first.

Source code in sleap_nn/architectures/heads.py

class MultiScaleHead(Head):
    """Base class for heads that read several backbone strides at once.

    ``Model`` binds a ``MultiScaleHead`` to the feature map at every stride in
    ``output_strides`` (decoder outputs, plus the bottleneck for the deepest one)
    and calls its layer with the list of those maps, finest first.
    ``output_stride`` is the finest of them.

    Attributes:
        output_strides: Feature strides to read, finest first.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(self, output_strides: Sequence[int], loss_weight: float = 1.0):
        """Initialize the object with the specified attributes."""
        self.output_strides = sorted(int(s) for s in output_strides)
        super().__init__(self.output_strides[0], loss_weight)

    def make_head(self, x_in: Sequence[int]) -> nn.Module:
        """Make the head from the input channels of each stride, finest first."""
        message = "Subclasses must implement this method."
        logger.error(message)
        raise NotImplementedError(message)

__init__(output_strides, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(self, output_strides: Sequence[int], loss_weight: float = 1.0):
    """Initialize the object with the specified attributes."""
    self.output_strides = sorted(int(s) for s in output_strides)
    super().__init__(self.output_strides[0], loss_weight)

make_head(x_in)

Make the head from the input channels of each stride, finest first.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: Sequence[int]) -> nn.Module:
    """Make the head from the input channels of each stride, finest first."""
    message = "Subclasses must implement this method."
    logger.error(message)
    raise NotImplementedError(message)

OffsetRefinementHead

Bases: Head

Head for specifying offset refinement maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma_threshold		Threshold of confidence map values to use for defining the boundary of the offset maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class OffsetRefinementHead(Head):
    """Head for specifying offset refinement maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma_threshold: Threshold of confidence map values to use for defining the
            boundary of the offset maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma_threshold: float = 0.2,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma_threshold = sigma_threshold

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return int(len(self.part_names) * 2)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
        sigma_threshold: float = 0.2,
        loss_weight: float = 1.0,
    ) -> "OffsetRefinementHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            part_names: Text name of the body parts (nodes) that the head will be
                configured to produce. The number of parts determines the number of
                channels in the output. This must be provided if the `part_names`
                attribute of the configuration is not set.
            sigma_threshold: Minimum confidence map value below which offsets will be
                replaced with zeros.
            loss_weight: Weight of the loss associated with this head.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if hasattr(config, "part_names"):
            if config.part_names is not None:
                part_names = config.part_names
        elif hasattr(config, "anchor_part"):
            part_names = [config.anchor_part]
        else:
            message = "Required attribute 'part_names' is missing in the configuration."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            output_stride=config.output_stride,
            sigma_threshold=sigma_threshold,
            loss_weight=loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma_threshold=0.2, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma_threshold: float = 0.2,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma_threshold = sigma_threshold

from_config(config, part_names=None, sigma_threshold=0.2, loss_weight=1.0) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		Text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. This must be provided if the part_names attribute of the configuration is not set.
sigma_threshold		Minimum confidence map value below which offsets will be replaced with zeros.
loss_weight		Weight of the loss associated with this head.

Returns:

Type	Description
OffsetRefinementHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
    sigma_threshold: float = 0.2,
    loss_weight: float = 1.0,
) -> "OffsetRefinementHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        part_names: Text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of
            channels in the output. This must be provided if the `part_names`
            attribute of the configuration is not set.
        sigma_threshold: Minimum confidence map value below which offsets will be
            replaced with zeros.
        loss_weight: Weight of the loss associated with this head.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if hasattr(config, "part_names"):
        if config.part_names is not None:
            part_names = config.part_names
    elif hasattr(config, "anchor_part"):
        part_names = [config.anchor_part]
    else:
        message = "Required attribute 'part_names' is missing in the configuration."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        output_stride=config.output_stride,
        sigma_threshold=sigma_threshold,
        loss_weight=loss_weight,
    )

PartAffinityFieldsHead

Bases: Head

Head for specifying multi-instance part affinity fields.

Attributes:

Name	Type	Description
edges		List of tuples of (source, destination) node names.
sigma		Spread of the part affinity fields.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class PartAffinityFieldsHead(Head):
    """Head for specifying multi-instance part affinity fields.

    Attributes:
        edges: List of tuples of `(source, destination)` node names.
        sigma: Spread of the part affinity fields.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        edges: Sequence[Tuple[Text, Text]],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.edges = edges
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return int(len(self.edges) * 2)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        edges: Optional[Sequence[Tuple[Text, Text]]] = None,
    ) -> "PartAffinityFieldsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            edges: List of 2-tuples of the form `(source_node, destination_node)` that
                define pairs of text names of the directed edges of the graph. This must
                be set if the `edges` attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.edges is not None:
            edges = config.edges
        return cls(
            edges=edges,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(edges, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    edges: Sequence[Tuple[Text, Text]],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.edges = edges
    self.sigma = sigma

from_config(config, edges=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
edges		List of 2-tuples of the form (source_node, destination_node) that define pairs of text names of the directed edges of the graph. This must be set if the edges attribute of the configuration is not set.

Returns:

Type	Description
PartAffinityFieldsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    edges: Optional[Sequence[Tuple[Text, Text]]] = None,
) -> "PartAffinityFieldsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        edges: List of 2-tuples of the form `(source_node, destination_node)` that
            define pairs of text names of the directed edges of the graph. This must
            be set if the `edges` attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.edges is not None:
        edges = config.edges
    return cls(
        edges=edges,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

SSDHead

Bases: MultiScaleHead

Multi-scale anchor head of a bbox_ssd (SSD-style) detector.

At every pyramid level, a small conv tower feeds two 3x3 predictors: per anchor, len(categories) + 1 class logits (index 0 is background) and 4 box deltas. The layer outputs one (B, A_total, len(categories) + 5) tensor with all levels' anchors concatenated (finest level first).

Attributes:

Name	Type	Description
categories		Ordered class names (background excluded).
anchors_per_location		Anchors at each feature-map cell.
output_strides		Pyramid level strides, finest first.
head_channels		Channels of each level's conv tower.
head_convs		3x3 conv + ReLU layers per tower.
loss_weight		Weight of the box regression loss.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
make_head	Make the per-level predictor from each level's input channels.

Source code in sleap_nn/architectures/heads.py

class SSDHead(MultiScaleHead):
    """Multi-scale anchor head of a ``bbox_ssd`` (SSD-style) detector.

    At every pyramid level, a small conv tower feeds two 3x3 predictors: per anchor,
    ``len(categories) + 1`` class logits (index 0 is background) and 4 box deltas.
    The layer outputs one ``(B, A_total, len(categories) + 5)`` tensor with all
    levels' anchors concatenated (finest level first).

    Attributes:
        categories: Ordered class names (background excluded).
        anchors_per_location: Anchors at each feature-map cell.
        output_strides: Pyramid level strides, finest first.
        head_channels: Channels of each level's conv tower.
        head_convs: 3x3 conv + ReLU layers per tower.
        loss_weight: Weight of the box regression loss.
    """

    def __init__(
        self,
        categories: Sequence[Text],
        anchors_per_location: int,
        output_strides: Sequence[int],
        head_channels: int = 128,
        head_convs: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_strides, loss_weight)
        self.categories = list(categories)
        self.anchors_per_location = int(anchors_per_location)
        self.head_channels = int(head_channels)
        self.head_convs = int(head_convs)

    @property
    def num_logits(self) -> int:
        """Class logits per anchor (background included)."""
        return len(self.categories) + 1

    @property
    def channels(self) -> int:
        """Values predicted per anchor (class logits + 4 box deltas)."""
        return self.num_logits + 4

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "ssd_multibox"

    def make_head(self, x_in: Sequence[int]) -> nn.Module:
        """Make the per-level predictor from each level's input channels."""
        return _SSDPredictor(
            in_channels=list(x_in),
            anchors_per_location=self.anchors_per_location,
            num_logits=self.num_logits,
            head_channels=self.head_channels,
            head_convs=self.head_convs,
        )

channels property

Values predicted per anchor (class logits + 4 box deltas).

loss_function property

Return the name of the loss function to use for this head.

num_logits property

Class logits per anchor (background included).

__init__(categories, anchors_per_location, output_strides, head_channels=128, head_convs=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    categories: Sequence[Text],
    anchors_per_location: int,
    output_strides: Sequence[int],
    head_channels: int = 128,
    head_convs: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_strides, loss_weight)
    self.categories = list(categories)
    self.anchors_per_location = int(anchors_per_location)
    self.head_channels = int(head_channels)
    self.head_convs = int(head_convs)

make_head(x_in)

Make the per-level predictor from each level's input channels.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: Sequence[int]) -> nn.Module:
    """Make the per-level predictor from each level's input channels."""
    return _SSDPredictor(
        in_channels=list(x_in),
        anchors_per_location=self.anchors_per_location,
        num_logits=self.num_logits,
        head_channels=self.head_channels,
        head_convs=self.head_convs,
    )

SegmentationHead

Bases: Head

Head for predicting binary foreground segmentation masks.

Outputs a single-channel map with sigmoid activation representing the probability that each pixel belongs to any instance (foreground).

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class SegmentationHead(Head):
    """Head for predicting binary foreground segmentation masks.

    Outputs a single-channel map with sigmoid activation representing the
    probability that each pixel belongs to any instance (foreground).

    Attributes:
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        output_stride: int = 2,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "bce_dice"

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=2, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    output_stride: int = 2,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

SingleInstanceConfmapsHead

Bases: Head

Head for specifying single instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class SingleInstanceConfmapsHead(Head):
    """Head for specifying single instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "SingleInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: Text name of the body parts (nodes) that the head will be
                configured to produce. The number of parts determines the number of
                channels in the output. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		Text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
SingleInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "SingleInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: Text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of
            channels in the output. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

default_embedding_pool(encoder_frozen)

The pooling an embedding head uses when its config leaves pool unset.

avg for a frozen pretrained encoder, gem for everything else. GeM floors its input at eps, so it pools only the positive part of each channel; a pretrained encoder's bottleneck is LayerNorm output (DINOv2-with-registers-base 50% negative, ConvNeXtV2-nano 55%), and a frozen one cannot adapt to that. Measured on the gerbil re-ID set (8 paired seeds, 3 epochs, val rank-1, avg - gem): frozen DINOv2-reg-base +0.042 (⅞ seeds, paired t p=0.007), frozen ConvNeXtV2-nano +0.046 (⅞, p=0.015); fine-tuned ConvNeXtV2-nano -0.001 (p=0.87) and fine-tuned DINOv2-reg-base -0.035 (p=0.31; one avg seed trained to 0.65 against GeM's 0.91) -- so fine-tuned pretrained encoders keep GeM, as do the native backbones, whose bottleneck ends in a ReLU.

Parameters:

Name	Type	Description	Default
encoder_frozen	bool	embedding_encoder_is_frozen(...) for the model.	required

Returns:

Type	Description
str	"avg" or "gem".

Source code in sleap_nn/architectures/heads.py

def default_embedding_pool(encoder_frozen: bool) -> str:
    """The pooling an `embedding` head uses when its config leaves ``pool`` unset.

    ``avg`` for a frozen pretrained encoder, ``gem`` for everything else. GeM floors
    its input at ``eps``, so it pools only the positive part of each channel; a
    pretrained encoder's bottleneck is LayerNorm output (DINOv2-with-registers-base
    50% negative, ConvNeXtV2-nano 55%), and a frozen one cannot adapt to that.
    Measured on the gerbil re-ID set (8 paired seeds, 3 epochs, val rank-1,
    avg - gem): frozen DINOv2-reg-base +0.042 (7/8 seeds, paired t p=0.007), frozen
    ConvNeXtV2-nano +0.046 (7/8, p=0.015); fine-tuned ConvNeXtV2-nano -0.001
    (p=0.87) and fine-tuned DINOv2-reg-base -0.035 (p=0.31; one avg seed trained
    to 0.65 against GeM's 0.91) -- so fine-tuned pretrained encoders keep GeM, as
    do the native backbones, whose bottleneck ends in a ReLU.

    Args:
        encoder_frozen: ``embedding_encoder_is_frozen(...)`` for the model.

    Returns:
        ``"avg"`` or ``"gem"``.
    """
    return "avg" if encoder_frozen else "gem"

embedding_encoder_is_frozen(backbone_type, backbone_config, embedding_config)

Whether an embedding model's pooled feature comes from a frozen pretrained encoder.

Parameters:

Name	Type	Description	Default
backbone_type	str	The backbone type (unet / convnext / swint / pretrained).	required
backbone_config		That backbone's config leaf.	required
embedding_config		The head_configs.embedding.embedding leaf.	required

Returns:

Type	Description
bool	True for a pretrained (HuggingFace) backbone frozen by backbone_config.pretrained.freeze or by the head's freeze_backbone. The native backbones always return False: their pooled feature is the output of randomly initialized, trainable middle blocks ending in a ReLU, frozen encoder or not.

Source code in sleap_nn/architectures/heads.py

def embedding_encoder_is_frozen(
    backbone_type: str, backbone_config, embedding_config
) -> bool:
    """Whether an `embedding` model's pooled feature comes from a frozen pretrained encoder.

    Args:
        backbone_type: The backbone type (``unet`` / ``convnext`` / ``swint`` /
            ``pretrained``).
        backbone_config: That backbone's config leaf.
        embedding_config: The ``head_configs.embedding.embedding`` leaf.

    Returns:
        ``True`` for a ``pretrained`` (HuggingFace) backbone frozen by
        ``backbone_config.pretrained.freeze`` or by the head's ``freeze_backbone``.
        The native backbones always return ``False``: their pooled feature is the
        output of randomly initialized, trainable middle blocks ending in a ReLU,
        frozen encoder or not.
    """
    from omegaconf import OmegaConf

    if backbone_type != "pretrained":
        return False
    return bool(OmegaConf.select(backbone_config, "freeze", default=False)) or bool(
        OmegaConf.select(embedding_config, "freeze_backbone", default=False)
    )