Skip to content

vit_neck

sleap_nn.architectures.vit_neck

Feature-pyramid neck that lets a plain (isotropic) ViT feed spatial heads.

A plain ViT -- DINOv2/v3, CLIP, SigLIP, MAE, ... -- emits every layer at a single stride (its patch size). The sleap-nn decoder needs a multi-scale pyramid to upsample through, so this neck turns one or more single-stride token maps into four levels at strides s/4, s/2, s, 2s (s = the token stride), in the style of ViTDet's Simple Feature Pyramid (Li et al. 2022) and DPT's reassemble stage (Ranftl et al. 2021):

  • "dpt" -- one tapped transformer layer per level, shallowest to the finest level (the default: intermediate layers keep more spatial detail).
  • "sfp" -- every level is built from the last tapped layer alone (ViTDet).

The pyramid then feeds the existing :class:~sleap_nn.architectures.encoder_decoder.Decoder exactly like the stage outputs of a hierarchical backbone, so heads need no changes.

Classes:

Name Description
DetailStem

Light convolutional stream giving a plain ViT high-resolution skips.
LayerNorm2d

LayerNorm over the channels of an NCHW map (torchvision's, as ConvNeXt uses).
ViTPyramidNeck

Build a feature pyramid from single-stride ViT feature maps.

DetailStem

Bases: Module

Light convolutional stream giving a plain ViT high-resolution skips.

A plain ViT's finest real features sit at its token stride (16), so the pyramid levels the neck synthesizes below it carry no image detail of their own. This stem (in the spirit of ViTMatte's detail-capture stream, Yao et al. 2024) reads the image directly and emits features at strides 2, 4 and 8, which are concatenated with the neck levels at the same stride (or added as new, finer levels) before the decoder.

Parameters:

Name Type Description Default
in_channels int

Input image channels.

 3
channels List[int]

Output channels at strides 2, 4, 8 (in that order).

 (32, 64, 128)

Methods:

Name Description
__init__

Initialize the stem.
forward

Return features at strides 2, 4, 8 (finest first).
Source code in sleap_nn/architectures/vit_neck.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
class DetailStem(nn.Module):
    """Light convolutional stream giving a plain ViT high-resolution skips.

    A plain ViT's finest real features sit at its token stride (16), so the
    pyramid levels the neck synthesizes below it carry no image detail of their
    own. This stem (in the spirit of ViTMatte's detail-capture stream, Yao et al.
    2024) reads the image directly and emits features at strides 2, 4 and 8,
    which are concatenated with the neck levels at the same stride (or added as
    new, finer levels) before the decoder.

    Args:
        in_channels: Input image channels.
        channels: Output channels at strides 2, 4, 8 (in that order).
    """

    def __init__(self, in_channels: int = 3, channels: List[int] = (32, 64, 128)):
        """Initialize the stem."""
        super().__init__()
        self.channels = list(channels)
        self.stages = nn.ModuleList()
        c_in = in_channels
        for c_out in self.channels:
            self.stages.append(
                nn.Sequential(
                    nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1),
                    LayerNorm2d(c_out),
                    nn.ReLU(inplace=True),
                    nn.Conv2d(c_out, c_out, kernel_size=3, padding=1),
                    LayerNorm2d(c_out),
                    nn.ReLU(inplace=True),
                )
            )
            c_in = c_out

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        """Return features at strides 2, 4, 8 (finest first)."""
        feats = []
        for stage in self.stages:
            x = stage(x)
            feats.append(x)
        return feats

__init__(in_channels=3, channels=(32, 64, 128))

Initialize the stem.

Source code in sleap_nn/architectures/vit_neck.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def __init__(self, in_channels: int = 3, channels: List[int] = (32, 64, 128)):
    """Initialize the stem."""
    super().__init__()
    self.channels = list(channels)
    self.stages = nn.ModuleList()
    c_in = in_channels
    for c_out in self.channels:
        self.stages.append(
            nn.Sequential(
                nn.Conv2d(c_in, c_out, kernel_size=3, stride=2, padding=1),
                LayerNorm2d(c_out),
                nn.ReLU(inplace=True),
                nn.Conv2d(c_out, c_out, kernel_size=3, padding=1),
                LayerNorm2d(c_out),
                nn.ReLU(inplace=True),
            )
        )
        c_in = c_out

forward(x)

Return features at strides 2, 4, 8 (finest first).

Source code in sleap_nn/architectures/vit_neck.py 
181
182
183
184
185
186
187
def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
    """Return features at strides 2, 4, 8 (finest first)."""
    feats = []
    for stage in self.stages:
        x = stage(x)
        feats.append(x)
    return feats

LayerNorm2d

Bases: LayerNorm2d

LayerNorm over the channels of an NCHW map (torchvision's, as ConvNeXt uses).

Checkpoints saved when this was a wrapper around an inner nn.LayerNorm store its parameters as <prefix>norm.weight / norm.bias; they are renamed on load so those models keep loading.

Methods:

Name Description
__init__

Initialize the norm.
Source code in sleap_nn/architectures/vit_neck.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class LayerNorm2d(_TorchvisionLayerNorm2d):
    """LayerNorm over the channels of an NCHW map (torchvision's, as ConvNeXt uses).

    Checkpoints saved when this was a wrapper around an inner ``nn.LayerNorm``
    store its parameters as ``<prefix>norm.weight`` / ``norm.bias``; they are
    renamed on load so those models keep loading.
    """

    def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
        """Initialize the norm."""
        super().__init__(num_channels, eps=eps)

    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
        for name in ("weight", "bias"):
            old = f"{prefix}norm.{name}"
            if old in state_dict:
                state_dict[prefix + name] = state_dict.pop(old)
        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)

__init__(num_channels, eps=1e-06)

Initialize the norm.

Source code in sleap_nn/architectures/vit_neck.py 
37
38
39
def __init__(self, num_channels: int, eps: float = 1e-6) -> None:
    """Initialize the norm."""
    super().__init__(num_channels, eps=eps)

ViTPyramidNeck

Bases: Module

Build a feature pyramid from single-stride ViT feature maps.

Parameters:

Name Type Description Default
in_channels int

Channel count of the ViT feature maps (its hidden size).

 required
out_channels List[int]

Channels of each output level, finest first (one per scale).

 required
num_inputs int

Number of tapped layers fed to forward. "dpt" needs one per level; "sfp" uses only the last one.

 4
mode str

"dpt" (one tapped layer per level) or "sfp" (the last tapped layer for every level).

 'dpt'
scales Sequence[float]

The levels to build, as scale factors of the token grid, finest first -- a suffix of PYRAMID_SCALES. Levels finer than the model's output stride are left out: the decoder would never read them, and unread levels are parameters without a gradient (a DDP error).

 PYRAMID_SCALES

Methods:

Name Description
__init__

Initialize the neck.
forward

Map tapped token maps to the pyramid, finest level first.
Source code in sleap_nn/architectures/vit_neck.py 
 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
class ViTPyramidNeck(nn.Module):
    """Build a feature pyramid from single-stride ViT feature maps.

    Args:
        in_channels: Channel count of the ViT feature maps (its hidden size).
        out_channels: Channels of each output level, finest first (one per scale).
        num_inputs: Number of tapped layers fed to `forward`. ``"dpt"`` needs one
            per level; ``"sfp"`` uses only the last one.
        mode: ``"dpt"`` (one tapped layer per level) or ``"sfp"`` (the last
            tapped layer for every level).
        scales: The levels to build, as scale factors of the token grid, finest
            first -- a suffix of `PYRAMID_SCALES`. Levels finer than the model's
            output stride are left out: the decoder would never read them, and
            unread levels are parameters without a gradient (a DDP error).
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: List[int],
        num_inputs: int = 4,
        mode: str = "dpt",
        scales: Sequence[float] = PYRAMID_SCALES,
    ) -> None:
        """Initialize the neck."""
        super().__init__()
        scales = [float(sc) for sc in scales]
        if mode not in ("dpt", "sfp"):
            raise ValueError(f"Unknown ViT neck mode {mode!r}; use 'dpt' or 'sfp'.")
        if len(out_channels) != len(scales):
            raise ValueError(
                f"ViT neck needs {len(scales)} output channel counts (one per level "
                f"at scales {scales}), got {list(out_channels)}."
            )
        if mode == "dpt" and num_inputs != len(scales):
            raise ValueError(
                f"The 'dpt' ViT neck taps one layer per pyramid level, so it needs "
                f"{len(scales)} layers, got {num_inputs}."
            )
        self.mode = mode
        self.in_channels = in_channels
        self.out_channels = list(out_channels)
        self.scales = scales

        self.levels = nn.ModuleList()
        for scale, c_out in zip(self.scales, self.out_channels):
            c_mid = _resampled_channels(in_channels, scale)
            self.levels.append(
                nn.Sequential(
                    _resample(in_channels, scale),
                    nn.Conv2d(c_mid, c_out, kernel_size=1, bias=False),
                    LayerNorm2d(c_out),
                    nn.Conv2d(c_out, c_out, kernel_size=3, padding=1, bias=False),
                    LayerNorm2d(c_out),
                )
            )

    def forward(self, feats: List[torch.Tensor]) -> List[torch.Tensor]:
        """Map tapped token maps to the pyramid, finest level first.

        Args:
            feats: Tapped ViT feature maps ``(B, in_channels, h, w)``, shallowest
                first.

        Returns:
            One map per level, at relative scales `self.scales` of the token grid.
        """
        if self.mode == "sfp":
            inputs = [feats[-1]] * len(self.levels)
        else:
            inputs = feats
        return [level(f) for level, f in zip(self.levels, inputs)]

__init__(in_channels, out_channels, num_inputs=4, mode='dpt', scales=PYRAMID_SCALES)

Initialize the neck.

Source code in sleap_nn/architectures/vit_neck.py 
 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
def __init__(
    self,
    in_channels: int,
    out_channels: List[int],
    num_inputs: int = 4,
    mode: str = "dpt",
    scales: Sequence[float] = PYRAMID_SCALES,
) -> None:
    """Initialize the neck."""
    super().__init__()
    scales = [float(sc) for sc in scales]
    if mode not in ("dpt", "sfp"):
        raise ValueError(f"Unknown ViT neck mode {mode!r}; use 'dpt' or 'sfp'.")
    if len(out_channels) != len(scales):
        raise ValueError(
            f"ViT neck needs {len(scales)} output channel counts (one per level "
            f"at scales {scales}), got {list(out_channels)}."
        )
    if mode == "dpt" and num_inputs != len(scales):
        raise ValueError(
            f"The 'dpt' ViT neck taps one layer per pyramid level, so it needs "
            f"{len(scales)} layers, got {num_inputs}."
        )
    self.mode = mode
    self.in_channels = in_channels
    self.out_channels = list(out_channels)
    self.scales = scales

    self.levels = nn.ModuleList()
    for scale, c_out in zip(self.scales, self.out_channels):
        c_mid = _resampled_channels(in_channels, scale)
        self.levels.append(
            nn.Sequential(
                _resample(in_channels, scale),
                nn.Conv2d(c_mid, c_out, kernel_size=1, bias=False),
                LayerNorm2d(c_out),
                nn.Conv2d(c_out, c_out, kernel_size=3, padding=1, bias=False),
                LayerNorm2d(c_out),
            )
        )

forward(feats)

Map tapped token maps to the pyramid, finest level first.

Parameters:

Name Type Description Default
feats List[Tensor]

Tapped ViT feature maps (B, in_channels, h, w), shallowest first.

 required

Returns:

Type Description
List[Tensor]

One map per level, at relative scales self.scales of the token grid.
Source code in sleap_nn/architectures/vit_neck.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def forward(self, feats: List[torch.Tensor]) -> List[torch.Tensor]:
    """Map tapped token maps to the pyramid, finest level first.

    Args:
        feats: Tapped ViT feature maps ``(B, in_channels, h, w)``, shallowest
            first.

    Returns:
        One map per level, at relative scales `self.scales` of the token grid.
    """
    if self.mode == "sfp":
        inputs = [feats[-1]] * len(self.levels)
    else:
        inputs = feats
    return [level(f) for level, f in zip(self.levels, inputs)]