Skip to content

deeplabcut.pose_estimation_pytorch.models.backbones.base

Classes:

Name Description
BaseBackbone

Base Backbone class for pose estimation.
HuggingFaceWeightsMixin

Mixin for backbones where the pretrained weights are stored on HuggingFace.

BaseBackbone

Bases: ABC, Module

Base Backbone class for pose estimation.

Attributes:

Name Type Description
stride

the stride for the backbone
freeze_bn_weights

freeze weights of batch norm layers during training
freeze_bn_stats

freeze stats of batch norm layers during training

Methods:

Name Description
forward

Abstract method for the forward pass through the backbone.
freeze_batch_norm_layers

Freezes batch norm layers.
train

Sets the module in training or evaluation mode.
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/base.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class BaseBackbone(ABC, nn.Module):
    """Base Backbone class for pose estimation.

    Attributes:
        stride: the stride for the backbone
        freeze_bn_weights: freeze weights of batch norm layers during training
        freeze_bn_stats: freeze stats of batch norm layers during training
    """

    def __init__(
        self,
        stride: int | float,
        freeze_bn_weights: bool = True,
        freeze_bn_stats: bool = True,
    ):
        super().__init__()
        self.stride = stride
        self.freeze_bn_weights = freeze_bn_weights
        self.freeze_bn_stats = freeze_bn_stats

    @abstractmethod
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Abstract method for the forward pass through the backbone.

        Args:
            x: Input tensor of shape (batch_size, channels, height, width).

        Returns:
            a feature map for the input, of shape (batch_size, c', h', w')
        """
        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.BatchNorm2d):
                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) -> torch.Tensor

Abstract method for the forward pass through the backbone.

Parameters:

Name Type Description Default

x

Tensor

Input tensor of shape (batch_size, channels, height, width).

 required

Returns:

Type Description
Tensor

a feature map for the input, of shape (batch_size, c', h', w')
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/base.py 
47
48
49
50
51
52
53
54
55
56
57
@abstractmethod
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Abstract method for the forward pass through the backbone.

    Args:
        x: Input tensor of shape (batch_size, channels, height, width).

    Returns:
        a feature map for the input, of shape (batch_size, c', h', w')
    """
    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/backbones/base.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
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.BatchNorm2d):
            if self.freeze_bn_weights:
                module.weight.requires_grad = False
                module.bias.requires_grad = False
            if self.freeze_bn_stats:
                module.eval()

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/backbones/base.py 
76
77
78
79
80
81
82
83
84
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()

HuggingFaceWeightsMixin

Mixin for backbones where the pretrained weights are stored on HuggingFace.

Methods:

Name Description
download_weights

Downloads the backbone weights from the HuggingFace repo.
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/base.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
class HuggingFaceWeightsMixin:
    """Mixin for backbones where the pretrained weights are stored on HuggingFace."""

    def __init__(
        self,
        backbone_weight_folder: str | Path | None = None,
        repo_id: str = "DeepLabCut/DeepLabCut-Backbones",
        *args,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        if backbone_weight_folder is None:
            backbone_weight_folder = Path(__file__).parent / "pretrained_weights"
        else:
            backbone_weight_folder = Path(backbone_weight_folder).resolve()

        self.backbone_weight_folder = backbone_weight_folder
        self.repo_id = repo_id

    def download_weights(self, filename: str, force: bool = False) -> Path:
        """Downloads the backbone weights from the HuggingFace repo.

        Args:
            filename: The name of the model file to download in the repo.
            force: Whether to re-download the file if it already exists locally.

        Returns:
            The path to the model snapshot.
        """
        model_path = self.backbone_weight_folder / filename
        if model_path.exists():
            if not force:
                return model_path
            model_path.unlink()

        logging.info(f"Downloading the pre-trained backbone to {model_path}")
        self.backbone_weight_folder.mkdir(exist_ok=True, parents=False)
        output_path = Path(hf_hub_download(self.repo_id, filename, cache_dir=self.backbone_weight_folder))

        # resolve gets the actual path if the output path is a symlink
        output_path = output_path.resolve()
        # move to the target path
        output_path.rename(model_path)

        # delete downloaded artifacts
        uid, rid = self.repo_id.split("/")
        artifact_dir = self.backbone_weight_folder / f"models--{uid}--{rid}"
        if artifact_dir.exists():
            shutil.rmtree(artifact_dir)

        return model_path

download_weights 

download_weights(filename: str, force: bool = False) -> Path

Downloads the backbone weights from the HuggingFace repo.

Parameters:

Name Type Description Default

filename

str

The name of the model file to download in the repo.

 required

force

bool

Whether to re-download the file if it already exists locally.

 False

Returns:

Type Description
Path

The path to the model snapshot.
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/base.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def download_weights(self, filename: str, force: bool = False) -> Path:
    """Downloads the backbone weights from the HuggingFace repo.

    Args:
        filename: The name of the model file to download in the repo.
        force: Whether to re-download the file if it already exists locally.

    Returns:
        The path to the model snapshot.
    """
    model_path = self.backbone_weight_folder / filename
    if model_path.exists():
        if not force:
            return model_path
        model_path.unlink()

    logging.info(f"Downloading the pre-trained backbone to {model_path}")
    self.backbone_weight_folder.mkdir(exist_ok=True, parents=False)
    output_path = Path(hf_hub_download(self.repo_id, filename, cache_dir=self.backbone_weight_folder))

    # resolve gets the actual path if the output path is a symlink
    output_path = output_path.resolve()
    # move to the target path
    output_path.rename(model_path)

    # delete downloaded artifacts
    uid, rid = self.repo_id.split("/")
    artifact_dir = self.backbone_weight_folder / f"models--{uid}--{rid}"
    if artifact_dir.exists():
        shutil.rmtree(artifact_dir)

    return model_path