Skip to content

deeplabcut.pose_estimation_pytorch.models.necks.layers

Classes:

Name Description
Attention

Attention block module.
FeedForward

FeedForward block module.
PreNorm

PreNorm block module.
Residual

Residual block module.
TransformerLayer

TransformerLayer block module.

Attention

Bases: Module

Attention block module.

This module implements the attention mechanism in the transformer layers.

Attributes:

Name Type Description
dim

Dimension of the input tensor.
heads

Number of attention heads. Defaults to 8.
dropout

Dropout rate. Defaults to 0.0.
num_keypoints

Number of keypoints. Defaults to None.
scale_with_head

Scale attention with the number of heads. Defaults to False.

Methods:

Name Description
__init__

Initialize the Attention block.
forward

Forward pass through the Attention block.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
class Attention(torch.nn.Module):
    """Attention block module.

    This module implements the attention mechanism in the transformer layers.

    Attributes:
        dim: Dimension of the input tensor.
        heads: Number of attention heads. Defaults to 8.
        dropout: Dropout rate. Defaults to 0.0.
        num_keypoints: Number of keypoints. Defaults to None.
        scale_with_head: Scale attention with the number of heads.
                         Defaults to False.
    """

    def __init__(
        self,
        dim: int,
        heads: int = 8,
        dropout: float = 0.0,
        num_keypoints: int = None,
        scale_with_head: bool = False,
    ):
        """Initialize the Attention block.

        Args:
            dim: Dimension of the input tensor.
            heads: Number of attention heads. Defaults to 8.
            dropout: Dropout rate. Defaults to 0.0.
            num_keypoints: Number of keypoints. Defaults to None.
            scale_with_head: Scale attention with the number of heads.
                             Defaults to False.
        """
        super().__init__()
        self.heads = heads
        self.scale = (dim // heads) ** -0.5 if scale_with_head else dim**-0.5

        self.to_qkv = torch.nn.Linear(dim, dim * 3, bias=False)
        self.to_out = torch.nn.Sequential(torch.nn.Linear(dim, dim), torch.nn.Dropout(dropout))
        self.num_keypoints = num_keypoints

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
        """Forward pass through the Attention block.

        Args:
            x: Input tensor.
            mask: Attention mask. Defaults to None.

        Returns:
            Output tensor.
        """
        _b, _n, _, h = *x.shape, self.heads
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)

        dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale
        mask_value = -torch.finfo(dots.dtype).max

        if mask is not None:
            mask = F.pad(mask.flatten(1), (1, 0), value=True)
            assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions"
            mask = mask[:, None, :] * mask[:, :, None]
            dots.masked_fill_(~mask, mask_value)
            del mask

        attn = dots.softmax(dim=-1)

        out = torch.einsum("bhij,bhjd->bhid", attn, v)

        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.to_out(out)
        return out

__init__ 

__init__(dim: int, heads: int = 8, dropout: float = 0.0, num_keypoints: int = None, scale_with_head: bool = False)

Initialize the Attention block.

Parameters:

Name Type Description Default

dim

int

Dimension of the input tensor.

 required

heads

int

Number of attention heads. Defaults to 8.

 8

dropout

float

Dropout rate. Defaults to 0.0.

 0.0

num_keypoints

int

Number of keypoints. Defaults to None.

 None

scale_with_head

bool

Scale attention with the number of heads. Defaults to False.

 False
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
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
def __init__(
    self,
    dim: int,
    heads: int = 8,
    dropout: float = 0.0,
    num_keypoints: int = None,
    scale_with_head: bool = False,
):
    """Initialize the Attention block.

    Args:
        dim: Dimension of the input tensor.
        heads: Number of attention heads. Defaults to 8.
        dropout: Dropout rate. Defaults to 0.0.
        num_keypoints: Number of keypoints. Defaults to None.
        scale_with_head: Scale attention with the number of heads.
                         Defaults to False.
    """
    super().__init__()
    self.heads = heads
    self.scale = (dim // heads) ** -0.5 if scale_with_head else dim**-0.5

    self.to_qkv = torch.nn.Linear(dim, dim * 3, bias=False)
    self.to_out = torch.nn.Sequential(torch.nn.Linear(dim, dim), torch.nn.Dropout(dropout))
    self.num_keypoints = num_keypoints

forward 

forward(x: Tensor, mask: Tensor = None)

Forward pass through the Attention block.

Parameters:

Name Type Description Default

x

Tensor

Input tensor.

 required

mask

Tensor

Attention mask. Defaults to None.

 None

Returns:

Type Description

Output tensor.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
    """Forward pass through the Attention block.

    Args:
        x: Input tensor.
        mask: Attention mask. Defaults to None.

    Returns:
        Output tensor.
    """
    _b, _n, _, h = *x.shape, self.heads
    qkv = self.to_qkv(x).chunk(3, dim=-1)
    q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)

    dots = torch.einsum("bhid,bhjd->bhij", q, k) * self.scale
    mask_value = -torch.finfo(dots.dtype).max

    if mask is not None:
        mask = F.pad(mask.flatten(1), (1, 0), value=True)
        assert mask.shape[-1] == dots.shape[-1], "mask has incorrect dimensions"
        mask = mask[:, None, :] * mask[:, :, None]
        dots.masked_fill_(~mask, mask_value)
        del mask

    attn = dots.softmax(dim=-1)

    out = torch.einsum("bhij,bhjd->bhid", attn, v)

    out = rearrange(out, "b h n d -> b n (h d)")
    out = self.to_out(out)
    return out

FeedForward

Bases: Module

FeedForward block module.

This module implements the feedforward layer in the transformer layers.

Attributes:

Name Type Description
dim

Dimension of the input tensor.
hidden_dim

Dimension of the hidden layer.
dropout

Dropout rate. Defaults to 0.0.

Methods:

Name Description
__init__

Initialize the FeedForward block.
forward

Forward pass through the FeedForward block.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
 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
class FeedForward(torch.nn.Module):
    """FeedForward block module.

    This module implements the feedforward layer in the transformer layers.

    Attributes:
        dim: Dimension of the input tensor.
        hidden_dim: Dimension of the hidden layer.
        dropout: Dropout rate. Defaults to 0.0.
    """

    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        """Initialize the FeedForward block.

        Args:
            dim: Dimension of the input tensor.
            hidden_dim: Dimension of the hidden layer.
            dropout: Dropout rate. Defaults to 0.0.
        """
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(dim, hidden_dim),
            torch.nn.GELU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(hidden_dim, dim),
            torch.nn.Dropout(dropout),
        )

    def forward(self, x: torch.Tensor):
        """Forward pass through the FeedForward block.

        Args:
            x: Input tensor.

        Returns:
            Output tensor.
        """
        return self.net(x)

__init__ 

__init__(dim: int, hidden_dim: int, dropout: float = 0.0)

Initialize the FeedForward block.

Parameters:

Name Type Description Default

dim

int

Dimension of the input tensor.

 required

hidden_dim

int

Dimension of the hidden layer.

 required

dropout

float

Dropout rate. Defaults to 0.0.

 0.0
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
    """Initialize the FeedForward block.

    Args:
        dim: Dimension of the input tensor.
        hidden_dim: Dimension of the hidden layer.
        dropout: Dropout rate. Defaults to 0.0.
    """
    super().__init__()
    self.net = torch.nn.Sequential(
        torch.nn.Linear(dim, hidden_dim),
        torch.nn.GELU(),
        torch.nn.Dropout(dropout),
        torch.nn.Linear(hidden_dim, dim),
        torch.nn.Dropout(dropout),
    )

forward 

forward(x: Tensor)

Forward pass through the FeedForward block.

Parameters:

Name Type Description Default

x

Tensor

Input tensor.

 required

Returns:

Type Description

Output tensor.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
114
115
116
117
118
119
120
121
122
123
def forward(self, x: torch.Tensor):
    """Forward pass through the FeedForward block.

    Args:
        x: Input tensor.

    Returns:
        Output tensor.
    """
    return self.net(x)

PreNorm

Bases: Module

PreNorm block module.

This module implements pre-normalization for the transformer layers.

Attributes:

Name Type Description
dim

Dimension of the input tensor.
fn

The function to apply after normalization.
fusion_factor

Fusion factor for layer normalization. Defaults to 1.

Methods:

Name Description
__init__

Initialize the PreNorm block.
forward

Forward pass through the PreNorm block.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
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
class PreNorm(torch.nn.Module):
    """PreNorm block module.

    This module implements pre-normalization for the transformer layers.

    Attributes:
        dim: Dimension of the input tensor.
        fn: The function to apply after normalization.
        fusion_factor: Fusion factor for layer normalization.
                       Defaults to 1.
    """

    def __init__(self, dim: int, fn: torch.nn.Module, fusion_factor: int = 1):
        """Initialize the PreNorm block.

        Args:
            dim: Dimension of the input tensor.
            fn: The function to apply after normalization.
            fusion_factor: Fusion factor for layer normalization.
                           Defaults to 1.
        """
        super().__init__()
        self.norm = torch.nn.LayerNorm(dim * fusion_factor)
        self.fn = fn

    def forward(self, x, **kwargs):
        """Forward pass through the PreNorm block.

        Args:
            x: Input tensor.
            **kwargs: Additional keyword arguments for the function.

        Returns:
            Output tensor.
        """
        return self.fn(self.norm(x), **kwargs)

__init__ 

__init__(dim: int, fn: Module, fusion_factor: int = 1)

Initialize the PreNorm block.

Parameters:

Name Type Description Default

dim

int

Dimension of the input tensor.

 required

fn

Module

The function to apply after normalization.

 required

fusion_factor

int

Fusion factor for layer normalization. Defaults to 1.

 1
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
60
61
62
63
64
65
66
67
68
69
70
71
def __init__(self, dim: int, fn: torch.nn.Module, fusion_factor: int = 1):
    """Initialize the PreNorm block.

    Args:
        dim: Dimension of the input tensor.
        fn: The function to apply after normalization.
        fusion_factor: Fusion factor for layer normalization.
                       Defaults to 1.
    """
    super().__init__()
    self.norm = torch.nn.LayerNorm(dim * fusion_factor)
    self.fn = fn

forward 

forward(x, **kwargs)

Forward pass through the PreNorm block.

Parameters:

Name Type Description Default

x

Input tensor.

 required

**kwargs

Additional keyword arguments for the function.

 {}

Returns:

Type Description

Output tensor.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
73
74
75
76
77
78
79
80
81
82
83
def forward(self, x, **kwargs):
    """Forward pass through the PreNorm block.

    Args:
        x: Input tensor.
        **kwargs: Additional keyword arguments for the function.

    Returns:
        Output tensor.
    """
    return self.fn(self.norm(x), **kwargs)

Residual

Bases: Module

Residual block module.

This module implements a residual block for the transformer layers.

Attributes:

Name Type Description
fn

The function to apply in the residual block.

Methods:

Name Description
__init__

Initialize the Residual block.
forward

Forward pass through the Residual block.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
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
class Residual(torch.nn.Module):
    """Residual block module.

    This module implements a residual block for the transformer layers.

    Attributes:
        fn: The function to apply in the residual block.
    """

    def __init__(self, fn: torch.nn.Module):
        """Initialize the Residual block.

        Args:
            fn: The function to apply in the residual block.
        """
        super().__init__()
        self.fn = fn

    def forward(self, x: torch.Tensor, **kwargs):
        """Forward pass through the Residual block.

        Args:
            x: Input tensor.
            **kwargs: Additional keyword arguments for the function.

        Returns:
            Output tensor.
        """
        return self.fn(x, **kwargs) + x

__init__ 

__init__(fn: Module)

Initialize the Residual block.

Parameters:

Name Type Description Default

fn

Module

The function to apply in the residual block.

 required
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
26
27
28
29
30
31
32
33
def __init__(self, fn: torch.nn.Module):
    """Initialize the Residual block.

    Args:
        fn: The function to apply in the residual block.
    """
    super().__init__()
    self.fn = fn

forward 

forward(x: Tensor, **kwargs)

Forward pass through the Residual block.

Parameters:

Name Type Description Default

x

Tensor

Input tensor.

 required

**kwargs

Additional keyword arguments for the function.

 {}

Returns:

Type Description

Output tensor.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
35
36
37
38
39
40
41
42
43
44
45
def forward(self, x: torch.Tensor, **kwargs):
    """Forward pass through the Residual block.

    Args:
        x: Input tensor.
        **kwargs: Additional keyword arguments for the function.

    Returns:
        Output tensor.
    """
    return self.fn(x, **kwargs) + x

TransformerLayer

Bases: Module

TransformerLayer block module.

This module implements the Transformer layer in the transformer model.

Attributes:

Name Type Description
dim

Dimension of the input tensor.
depth

Depth of the transformer layer.
heads

Number of attention heads.
mlp_dim

Dimension of the MLP layer.
dropout

Dropout rate.
num_keypoints

Number of keypoints. Defaults to None.
all_attn

Apply attention to all keypoints. Defaults to False.
scale_with_head

Scale attention with the number of heads. Defaults to False.

Methods:

Name Description
__init__

Initialize the TransformerLayer block.
forward

Forward pass through the TransformerLayer block.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
class TransformerLayer(torch.nn.Module):
    """TransformerLayer block module.

    This module implements the Transformer layer in the transformer model.

    Attributes:
        dim: Dimension of the input tensor.
        depth: Depth of the transformer layer.
        heads: Number of attention heads.
        mlp_dim: Dimension of the MLP layer.
        dropout: Dropout rate.
        num_keypoints: Number of keypoints. Defaults to None.
        all_attn: Apply attention to all keypoints.
                  Defaults to False.
        scale_with_head: Scale attention with the number of heads.
                         Defaults to False.
    """

    def __init__(
        self,
        dim: int,
        depth: int,
        heads: int,
        mlp_dim: int,
        dropout: float,
        num_keypoints: int = None,
        all_attn: bool = False,
        scale_with_head: bool = False,
    ):
        """Initialize the TransformerLayer block.

        Args:
            dim: Dimension of the input tensor.
            depth: Depth of the transformer layer.
            heads: Number of attention heads.
            mlp_dim: Dimension of the MLP layer.
            dropout: Dropout rate.
            num_keypoints: Number of keypoints. Defaults to None.
            all_attn: Apply attention to all keypoints. Defaults to False.
            scale_with_head: Scale attention with the number of heads. Defaults to False.
        """
        super().__init__()
        self.layers = torch.nn.ModuleList([])
        self.all_attn = all_attn
        self.num_keypoints = num_keypoints
        for _ in range(depth):
            self.layers.append(
                torch.nn.ModuleList(
                    [
                        Residual(
                            PreNorm(
                                dim,
                                Attention(
                                    dim,
                                    heads=heads,
                                    dropout=dropout,
                                    num_keypoints=num_keypoints,
                                    scale_with_head=scale_with_head,
                                ),
                            )
                        ),
                        Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))),
                    ]
                )
            )

    def forward(self, x: torch.Tensor, mask: torch.Tensor = None, pos: torch.Tensor = None):
        """Forward pass through the TransformerLayer block.

        Args:
            x: Input tensor.
            mask: Attention mask. Defaults to None.
            pos: Positional encoding. Defaults to None.

        Returns:
            Output tensor.
        """
        for idx, (attn, ff) in enumerate(self.layers):
            if idx > 0 and self.all_attn:
                x[:, self.num_keypoints :] += pos
            x = attn(x, mask=mask)
            x = ff(x)
        return x

__init__ 

__init__(
    dim: int,
    depth: int,
    heads: int,
    mlp_dim: int,
    dropout: float,
    num_keypoints: int = None,
    all_attn: bool = False,
    scale_with_head: bool = False,
)

Initialize the TransformerLayer block.

Parameters:

Name Type Description Default

dim

int

Dimension of the input tensor.

 required

depth

int

Depth of the transformer layer.

 required

heads

int

Number of attention heads.

 required

mlp_dim

int

Dimension of the MLP layer.

 required

dropout

float

Dropout rate.

 required

num_keypoints

int

Number of keypoints. Defaults to None.

 None

all_attn

bool

Apply attention to all keypoints. Defaults to False.

 False

scale_with_head

bool

Scale attention with the number of heads. Defaults to False.

 False
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
def __init__(
    self,
    dim: int,
    depth: int,
    heads: int,
    mlp_dim: int,
    dropout: float,
    num_keypoints: int = None,
    all_attn: bool = False,
    scale_with_head: bool = False,
):
    """Initialize the TransformerLayer block.

    Args:
        dim: Dimension of the input tensor.
        depth: Depth of the transformer layer.
        heads: Number of attention heads.
        mlp_dim: Dimension of the MLP layer.
        dropout: Dropout rate.
        num_keypoints: Number of keypoints. Defaults to None.
        all_attn: Apply attention to all keypoints. Defaults to False.
        scale_with_head: Scale attention with the number of heads. Defaults to False.
    """
    super().__init__()
    self.layers = torch.nn.ModuleList([])
    self.all_attn = all_attn
    self.num_keypoints = num_keypoints
    for _ in range(depth):
        self.layers.append(
            torch.nn.ModuleList(
                [
                    Residual(
                        PreNorm(
                            dim,
                            Attention(
                                dim,
                                heads=heads,
                                dropout=dropout,
                                num_keypoints=num_keypoints,
                                scale_with_head=scale_with_head,
                            ),
                        )
                    ),
                    Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))),
                ]
            )
        )

forward 

forward(x: Tensor, mask: Tensor = None, pos: Tensor = None)

Forward pass through the TransformerLayer block.

Parameters:

Name Type Description Default

x

Tensor

Input tensor.

 required

mask

Tensor

Attention mask. Defaults to None.

 None

pos

Tensor

Positional encoding. Defaults to None.

 None

Returns:

Type Description

Output tensor.
Source code in deeplabcut/pose_estimation_pytorch/models/necks/layers.py 
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, pos: torch.Tensor = None):
    """Forward pass through the TransformerLayer block.

    Args:
        x: Input tensor.
        mask: Attention mask. Defaults to None.
        pos: Positional encoding. Defaults to None.

    Returns:
        Output tensor.
    """
    for idx, (attn, ff) in enumerate(self.layers):
        if idx > 0 and self.all_attn:
            x[:, self.num_keypoints :] += pos
        x = attn(x, mask=mask)
        x = ff(x)
    return x