Skip to content

common

sleap_nn.architectures.common

Common utilities for architecture and model building.

Classes:

Name Description
FreezableEncoderMixin

Freeze a backbone's encoder and keep it in eval mode for the whole run.
MaxPool2dWithSamePadding

A MaxPool2d module with support for same padding.

FreezableEncoderMixin

Freeze a backbone's encoder and keep it in eval mode for the whole run.

The encoder is the part of the backbone that pretrained weights load into: the HuggingFace model of a pretrained backbone, the ImageNet enc of a native convnext / swint backbone, the stem + encoder blocks of a unet. Middle blocks and decoders are randomly initialized on those backbones and are NOT part of it. Subclasses whose encoder is not self.enc override encoder_modules.

Freezing means two things, and requires_grad_(False) alone does only the first: (1) no gradient reaches the encoder weights, and (2) the encoder runs in eval mode, so BatchNorm normalizes with its stored running statistics (and stops updating them) and dropout / stochastic depth are off. A "frozen" encoder that stays in train mode still changes: its BatchNorm statistics drift toward the new data and every training forward normalizes with batch statistics, so the features at inference differ from the ones the head was trained on.

Place the mixin before nn.Module in the bases so its train wins.

Methods:

Name Description
encoder_modules

Return the modules that make up the encoder (default: [self.enc]).
freeze_encoder

Freeze the encoder: no gradient to its weights, eval mode from now on.
train

Set train/eval mode, keeping a FROZEN encoder in eval.
Source code in sleap_nn/architectures/common.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
class FreezableEncoderMixin:
    """Freeze a backbone's encoder and keep it in eval mode for the whole run.

    The encoder is the part of the backbone that pretrained weights load into: the
    HuggingFace model of a ``pretrained`` backbone, the ImageNet ``enc`` of a native
    ``convnext`` / ``swint`` backbone, the stem + encoder blocks of a ``unet``. Middle
    blocks and decoders are randomly initialized on those backbones and are NOT part
    of it. Subclasses whose encoder is not ``self.enc`` override ``encoder_modules``.

    Freezing means two things, and ``requires_grad_(False)`` alone does only the
    first: (1) no gradient reaches the encoder weights, and (2) the encoder runs in
    eval mode, so BatchNorm normalizes with its stored running statistics (and stops
    updating them) and dropout / stochastic depth are off. A "frozen" encoder that
    stays in train mode still changes: its BatchNorm statistics drift toward the new
    data and every training forward normalizes with batch statistics, so the
    features at inference differ from the ones the head was trained on.

    Place the mixin before ``nn.Module`` in the bases so its ``train`` wins.
    """

    #: Set by ``freeze_encoder``; read by ``train``.
    freeze: bool = False

    def encoder_modules(self) -> List[nn.Module]:
        """Return the modules that make up the encoder (default: ``[self.enc]``)."""
        return [self.enc]

    def freeze_encoder(self) -> None:
        """Freeze the encoder: no gradient to its weights, eval mode from now on."""
        self.freeze = True
        for module in self.encoder_modules():
            module.eval()
            module.requires_grad_(False)

    def train(self, mode: bool = True):
        """Set train/eval mode, keeping a FROZEN encoder in eval.

        ``freeze_encoder`` puts the encoder in eval once, but Lightning calls
        ``model.train()`` at the start of every training epoch and that recurses
        into every submodule, so without this override the encoder goes back into
        train mode and its normalization layers keep updating. Weights stay frozen
        (``requires_grad_(False)``); the running statistics did not, which makes a
        "frozen" run irreproducible in a way that looks like seed noise (measured on
        ``microsoft/resnet-18``: ``running_mean`` moved 0.63 in a single forward).

        Only backbones with running statistics (BatchNorm: ResNet, BiT) or
        train-only stochasticity (the native ``swint``'s stochastic depth) are
        affected; LayerNorm-only encoders (ConvNeXt, ViT) compute the same thing in
        either mode.
        """
        super().train(mode)
        if getattr(self, "freeze", False):
            for module in self.encoder_modules():
                module.eval()
        return self

encoder_modules()

Return the modules that make up the encoder (default: [self.enc]).

Source code in sleap_nn/architectures/common.py 
33
34
35
def encoder_modules(self) -> List[nn.Module]:
    """Return the modules that make up the encoder (default: ``[self.enc]``)."""
    return [self.enc]

freeze_encoder()

Freeze the encoder: no gradient to its weights, eval mode from now on.

Source code in sleap_nn/architectures/common.py 
37
38
39
40
41
42
def freeze_encoder(self) -> None:
    """Freeze the encoder: no gradient to its weights, eval mode from now on."""
    self.freeze = True
    for module in self.encoder_modules():
        module.eval()
        module.requires_grad_(False)

train(mode=True)

Set train/eval mode, keeping a FROZEN encoder in eval.

freeze_encoder puts the encoder in eval once, but Lightning calls model.train() at the start of every training epoch and that recurses into every submodule, so without this override the encoder goes back into train mode and its normalization layers keep updating. Weights stay frozen (requires_grad_(False)); the running statistics did not, which makes a "frozen" run irreproducible in a way that looks like seed noise (measured on microsoft/resnet-18: running_mean moved 0.63 in a single forward).

Only backbones with running statistics (BatchNorm: ResNet, BiT) or train-only stochasticity (the native swint's stochastic depth) are affected; LayerNorm-only encoders (ConvNeXt, ViT) compute the same thing in either mode.

Source code in sleap_nn/architectures/common.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def train(self, mode: bool = True):
    """Set train/eval mode, keeping a FROZEN encoder in eval.

    ``freeze_encoder`` puts the encoder in eval once, but Lightning calls
    ``model.train()`` at the start of every training epoch and that recurses
    into every submodule, so without this override the encoder goes back into
    train mode and its normalization layers keep updating. Weights stay frozen
    (``requires_grad_(False)``); the running statistics did not, which makes a
    "frozen" run irreproducible in a way that looks like seed noise (measured on
    ``microsoft/resnet-18``: ``running_mean`` moved 0.63 in a single forward).

    Only backbones with running statistics (BatchNorm: ResNet, BiT) or
    train-only stochasticity (the native ``swint``'s stochastic depth) are
    affected; LayerNorm-only encoders (ConvNeXt, ViT) compute the same thing in
    either mode.
    """
    super().train(mode)
    if getattr(self, "freeze", False):
        for module in self.encoder_modules():
            module.eval()
    return self

MaxPool2dWithSamePadding

Bases: MaxPool2d

A MaxPool2d module with support for same padding.

This class extends the torch.nn.MaxPool2d module and adds the ability to perform 'same' padding, similar to 'same' padding in convolutional layers. When 'same' padding is specified, the input tensor is padded with zeros to ensure that the output spatial dimensions match the input spatial dimensions as closely as possible.

Parameters:

Name Type Description Default
nn.MaxPool2d arguments

Arguments that are passed to the parent torch.nn.MaxPool2d class.

 required

Methods:

Name Description
forward

torch.Tensor) -> torch.Tensor: Forward pass through the MaxPool2dWithSamePadding module.
Note

The 'same' padding is applied only when self.padding is set to "same".

Example

Create an instance of MaxPool2dWithSamePadding

maxpool_layer = MaxPool2dWithSamePadding(kernel_size=3, stride=2, padding="same")

Perform a forward pass on an input tensor

input_tensor = torch.rand(1, 3, 32, 32) # Example input tensor output = maxpool_layer(input_tensor) # Apply the MaxPool2d operation with same padding.

Source code in sleap_nn/architectures/common.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
class MaxPool2dWithSamePadding(nn.MaxPool2d):
    """A MaxPool2d module with support for same padding.

    This class extends the torch.nn.MaxPool2d module and adds the ability
    to perform 'same' padding, similar to 'same' padding in convolutional
    layers. When 'same' padding is specified, the input tensor is padded
    with zeros to ensure that the output spatial dimensions match the input
    spatial dimensions as closely as possible.

    Args:
        nn.MaxPool2d arguments: Arguments that are passed to the parent
            torch.nn.MaxPool2d class.

    Methods:
        forward(x: torch.Tensor) -> torch.Tensor:
            Forward pass through the MaxPool2dWithSamePadding module.

    Note:
        The 'same' padding is applied only when self.padding is set to "same".

    Example:
        # Create an instance of MaxPool2dWithSamePadding
        maxpool_layer = MaxPool2dWithSamePadding(kernel_size=3, stride=2, padding="same")

        # Perform a forward pass on an input tensor
        input_tensor = torch.rand(1, 3, 32, 32)  # Example input tensor
        output = maxpool_layer(input_tensor)  # Apply the MaxPool2d operation with same padding.
    """

    def __init__(self, *args, **kwargs):
        """Initialize the MaxPool2dWithSamePadding module."""
        super().__init__(*args, **kwargs)

    def _calc_same_pad(self, i: int, k: int, s: int, d: int) -> int:
        """Calculate the required padding to achieve 'same' padding.

        Args:
            i (int): Input dimension (height or width).
            k (int): Kernel size.
            s (int): Stride.
            d (int): Dilation.

        Returns:
            int: The calculated padding value.

        Note:
            Integer arithmetic (``-(-i // s)`` is ``ceil(i / s)``), not a tensor
            + ``.item()``: the sizes are Python ints, and ``Tensor.item()`` makes
            ``torch.compile`` split the graph at every pooling layer (and fail
            with ``fullgraph=True`` on PyTorch 2.6). No float division either,
            so it stays exact and symbolic under ``torch.compile(dynamic=True)``.
        """
        return max((-(-i // s) - 1) * s + (k - 1) * d + 1 - i, 0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the MaxPool2dWithSamePadding module.

        Args:
            x (torch.Tensor): Input tensor.

        Returns:
            torch.Tensor: Output tensor after applying the MaxPool2d operation.
        """
        padding = self.padding
        if padding == "same":
            ih, iw = x.size()[-2:]

            pad_h = self._calc_same_pad(
                i=ih,
                k=(
                    self.kernel_size
                    if type(self.kernel_size) is int
                    else self.kernel_size[0]
                ),
                s=self.stride if type(self.stride) is int else self.stride[0],
                d=self.dilation if type(self.dilation) is int else self.dilation[0],
            )
            pad_w = self._calc_same_pad(
                i=iw,
                k=(
                    self.kernel_size
                    if type(self.kernel_size) is int
                    else self.kernel_size[1]
                ),
                s=self.stride if type(self.stride) is int else self.stride[1],
                d=self.dilation if type(self.dilation) is int else self.dilation[1],
            )

            if pad_h > 0 or pad_w > 0:
                x = F.pad(
                    x, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
                )
            # A local, not ``self.padding = 0``: mutating the module made "same"
            # padding apply on the first forward only.
            padding = 0

        return F.max_pool2d(
            x,
            self.kernel_size,
            self.stride,
            padding,
            self.dilation,
            ceil_mode=self.ceil_mode,
            return_indices=self.return_indices,
        )

__init__(*args, **kwargs)

Initialize the MaxPool2dWithSamePadding module.

Source code in sleap_nn/architectures/common.py 
96
97
98
def __init__(self, *args, **kwargs):
    """Initialize the MaxPool2dWithSamePadding module."""
    super().__init__(*args, **kwargs)

forward(x)

Forward pass through the MaxPool2dWithSamePadding module.

Parameters:

Name Type Description Default
x Tensor

Input tensor.

 required

Returns:

Type Description
Tensor

torch.Tensor: Output tensor after applying the MaxPool2d operation.
Source code in sleap_nn/architectures/common.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the MaxPool2dWithSamePadding module.

    Args:
        x (torch.Tensor): Input tensor.

    Returns:
        torch.Tensor: Output tensor after applying the MaxPool2d operation.
    """
    padding = self.padding
    if padding == "same":
        ih, iw = x.size()[-2:]

        pad_h = self._calc_same_pad(
            i=ih,
            k=(
                self.kernel_size
                if type(self.kernel_size) is int
                else self.kernel_size[0]
            ),
            s=self.stride if type(self.stride) is int else self.stride[0],
            d=self.dilation if type(self.dilation) is int else self.dilation[0],
        )
        pad_w = self._calc_same_pad(
            i=iw,
            k=(
                self.kernel_size
                if type(self.kernel_size) is int
                else self.kernel_size[1]
            ),
            s=self.stride if type(self.stride) is int else self.stride[1],
            d=self.dilation if type(self.dilation) is int else self.dilation[1],
        )

        if pad_h > 0 or pad_w > 0:
            x = F.pad(
                x, (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
            )
        # A local, not ``self.padding = 0``: mutating the module made "same"
        # padding apply on the first forward only.
        padding = 0

    return F.max_pool2d(
        x,
        self.kernel_size,
        self.stride,
        padding,
        self.dilation,
        ceil_mode=self.ceil_mode,
        return_indices=self.return_indices,
    )