Skip to content

deeplabcut.pose_estimation_pytorch.models.detectors.base

Classes:

Name Description
BaseDetector

Definition of the class BaseDetector object.

BaseDetector

Bases: ABC, Module

Definition of the class BaseDetector object.

This is an abstract class defining the common structure and inference for detectors.

Methods:

Name Description
forward

Forward pass of the detector.

freeze_batch_norm_layers

Freezes batch norm layers.

get_target

Get the target for training the detector.

train

Sets the module in training or evaluation mode.

Source code in deeplabcut/pose_estimation_pytorch/models/detectors/base.py
class BaseDetector(ABC, nn.Module):
    """Definition of the class BaseDetector object.

    This is an abstract class defining the common structure and inference for detectors.
    """

    def __init__(
        self,
        freeze_bn_stats: bool = False,
        freeze_bn_weights: bool = False,
        pretrained: bool = False,
    ) -> None:
        super().__init__()
        self.freeze_bn_stats = freeze_bn_stats
        self.freeze_bn_weights = freeze_bn_weights
        self._pretrained = pretrained

    @abstractmethod
    def forward(
        self, x: torch.Tensor, targets: list[dict[str, torch.Tensor]] | None = None
    ) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]:
        """Forward pass of the detector.

        Args:
            x: images to be processed
            targets: ground-truth boxes present in each images

        Returns:
            losses: {'loss_name': loss_value}
            detections: for each of the b images, {"boxes": bounding_boxes}
        """
        pass

    @abstractmethod
    def get_target(self, labels: dict) -> list[dict]:
        """Get the target for training the detector.

        Args:
            labels: annotations containing keypoints, bounding boxes, etc.

        Returns:
            list of dictionaries, each representing target information for a single annotation.
        """
        pass

    def freeze_batch_norm_layers(self) -> None:
        """Freezes batch norm layers.

        Running mean + var are always given to F.batch_norm, except when the layer is
        in `train` mode and track_running_stats is False, see
            https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html
        So to 'freeze' the running stats, the only way is to set the layer to "eval"
        mode.
        """
        for module in self.modules():
            if isinstance(module, nn.modules.batchnorm._BatchNorm):
                if self.freeze_bn_weights:
                    module.weight.requires_grad = False
                    module.bias.requires_grad = False
                if self.freeze_bn_stats:
                    module.eval()

    def train(self, mode: bool = True) -> None:
        """Sets the module in training or evaluation mode.

        Args:
            mode: whether to set training mode (True) or evaluation mode (False)
        """
        super().train(mode)
        if self.freeze_bn_weights or self.freeze_bn_stats:
            self.freeze_batch_norm_layers()

forward abstractmethod

forward(
    x: Tensor, targets: list[dict[str, Tensor]] | None = None
) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]

Forward pass of the detector.

Parameters:

Name Type Description Default

x

Tensor

images to be processed

required

targets

list[dict[str, Tensor]] | None

ground-truth boxes present in each images

None

Returns:

Name Type Description
losses tuple[dict[str, Tensor], list[dict[str, Tensor]]]

{'loss_name': loss_value} detections: for each of the b images, {"boxes": bounding_boxes}

Source code in deeplabcut/pose_estimation_pytorch/models/detectors/base.py
@abstractmethod
def forward(
    self, x: torch.Tensor, targets: list[dict[str, torch.Tensor]] | None = None
) -> tuple[dict[str, torch.Tensor], list[dict[str, torch.Tensor]]]:
    """Forward pass of the detector.

    Args:
        x: images to be processed
        targets: ground-truth boxes present in each images

    Returns:
        losses: {'loss_name': loss_value}
        detections: for each of the b images, {"boxes": bounding_boxes}
    """
    pass

freeze_batch_norm_layers

freeze_batch_norm_layers() -> None

Freezes batch norm layers.

Running mean + var are always given to F.batch_norm, except when the layer is in train mode and track_running_stats is False, see https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html So to 'freeze' the running stats, the only way is to set the layer to "eval" mode.

Source code in deeplabcut/pose_estimation_pytorch/models/detectors/base.py
def freeze_batch_norm_layers(self) -> None:
    """Freezes batch norm layers.

    Running mean + var are always given to F.batch_norm, except when the layer is
    in `train` mode and track_running_stats is False, see
        https://pytorch.org/docs/stable/_modules/torch/nn/modules/batchnorm.html
    So to 'freeze' the running stats, the only way is to set the layer to "eval"
    mode.
    """
    for module in self.modules():
        if isinstance(module, nn.modules.batchnorm._BatchNorm):
            if self.freeze_bn_weights:
                module.weight.requires_grad = False
                module.bias.requires_grad = False
            if self.freeze_bn_stats:
                module.eval()

get_target abstractmethod

get_target(labels: dict) -> list[dict]

Get the target for training the detector.

Parameters:

Name Type Description Default

labels

dict

annotations containing keypoints, bounding boxes, etc.

required

Returns:

Type Description
list[dict]

list of dictionaries, each representing target information for a single annotation.

Source code in deeplabcut/pose_estimation_pytorch/models/detectors/base.py
@abstractmethod
def get_target(self, labels: dict) -> list[dict]:
    """Get the target for training the detector.

    Args:
        labels: annotations containing keypoints, bounding boxes, etc.

    Returns:
        list of dictionaries, each representing target information for a single annotation.
    """
    pass

train

train(mode: bool = True) -> None

Sets the module in training or evaluation mode.

Parameters:

Name Type Description Default

mode

bool

whether to set training mode (True) or evaluation mode (False)

True
Source code in deeplabcut/pose_estimation_pytorch/models/detectors/base.py
def train(self, mode: bool = True) -> None:
    """Sets the module in training or evaluation mode.

    Args:
        mode: whether to set training mode (True) or evaluation mode (False)
    """
    super().train(mode)
    if self.freeze_bn_weights or self.freeze_bn_stats:
        self.freeze_batch_norm_layers()