Skip to content

deeplabcut.pose_estimation_pytorch.models.modules.coam_module

Classes:

Name Description
CoAMBlock

Conditional Attention Module (CoAM) block.
ScaledDotProductAttention

Scaled dot-product attention.
SimplifiedScaledDotProductAttention

Scaled dot-product attention.

CoAMBlock

Bases: Module

Conditional Attention Module (CoAM) block.

Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
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
class CoAMBlock(nn.Module):
    """Conditional Attention Module (CoAM) block."""

    def __init__(self, spat_dims, channel_list, cond_enc, n_heads=1, channel_only=False):
        super().__init__()
        self.att_layers = []
        self.spat_dims = spat_dims
        self.cond_enc = cond_enc
        d_cond = cond_enc.num_channels
        for i in range(len(spat_dims)):
            att_layer = DAModule(
                d_model=channel_list[i],
                d_cond=d_cond,
                kernel_size=3,
                H=spat_dims[i][1],
                W=spat_dims[i][0],
                n_heads=n_heads,
                channel_only=channel_only,
            )
            self.att_layers.append(att_layer)
        self.att_layers = nn.ModuleList(self.att_layers)

    def forward(self, y_list, cond_hm):
        # if not isinstance(self.cond_enc, (StackedKeypointEncoder, ColoredKeypointEncoder)):
        #     cond_hm = cond_hm[:,0].unsqueeze(1) # we only want one channel of the heatmap
        y_list_att = []
        for i in range(len(y_list)):
            y_att = self.att_layers[i](
                y_list[i],
                TF.resize(cond_hm, (self.spat_dims[i][1], self.spat_dims[i][0])),
            )
            y_list_att.append(y_att)
        return y_list_att

ScaledDotProductAttention

Bases: Module

Scaled dot-product attention.

Methods:

Name Description
__init__

:param d_model: Output dimensionality of the model
forward

Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
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
class ScaledDotProductAttention(nn.Module):
    """Scaled dot-product attention."""

    def __init__(self, in_dim_q, in_dim_k, d_k, d_v, h, dropout=0.1, rev=False):
        """
        :param d_model: Output dimensionality of the model
        :param d_k: Dimensionality of queries and keys
        :param d_v: Dimensionality of values
        :param h: Number of heads
        """
        super().__init__()

        # 'rev': condition is key/value and orig. feature map is query
        if rev:
            d_model = in_dim_q
        else:
            d_model = in_dim_k
        self.fc_q = nn.Linear(in_dim_q, h * d_k)
        self.fc_k = nn.Linear(in_dim_k, h * d_k)
        self.fc_v = nn.Linear(in_dim_k, h * d_v)
        self.fc_o = nn.Linear(h * d_v, d_model)
        self.dropout = nn.Dropout(dropout)

        self.d_model = d_model
        self.d_k = d_k
        self.d_v = d_v
        self.h = h

        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode="fan_out")
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, queries, keys, values, attention_mask=None, attention_weights=None):
        """Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
        nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask:
        Mask over attention values (b_s, h, nq, nk).

        True indicates masking.
        :param attention_weights: Multiplicative weights for attention values (b_s, h,
            nq, nk).
        :return:
        """
        b_s, nq = queries.shape[:2]
        nk = keys.shape[1]

        q = self.fc_q(queries).view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)
        k = self.fc_k(keys).view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)
        v = self.fc_v(values).view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)

        att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, nk)
        if attention_weights is not None:
            att = att * attention_weights
        if attention_mask is not None:
            att = att.masked_fill(attention_mask, -np.inf)
        att = torch.softmax(att, -1)
        att = self.dropout(att)

        out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)
        out = self.fc_o(out)  # (b_s, nq, d_model)
        return out

__init__ 

__init__(in_dim_q, in_dim_k, d_k, d_v, h, dropout=0.1, rev=False)

:param d_model: Output dimensionality of the model :param d_k: Dimensionality of queries and keys :param d_v: Dimensionality of values :param h: Number of heads

Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
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
def __init__(self, in_dim_q, in_dim_k, d_k, d_v, h, dropout=0.1, rev=False):
    """
    :param d_model: Output dimensionality of the model
    :param d_k: Dimensionality of queries and keys
    :param d_v: Dimensionality of values
    :param h: Number of heads
    """
    super().__init__()

    # 'rev': condition is key/value and orig. feature map is query
    if rev:
        d_model = in_dim_q
    else:
        d_model = in_dim_k
    self.fc_q = nn.Linear(in_dim_q, h * d_k)
    self.fc_k = nn.Linear(in_dim_k, h * d_k)
    self.fc_v = nn.Linear(in_dim_k, h * d_v)
    self.fc_o = nn.Linear(h * d_v, d_model)
    self.dropout = nn.Dropout(dropout)

    self.d_model = d_model
    self.d_k = d_k
    self.d_v = d_v
    self.h = h

    self.init_weights()

forward 

forward(queries, keys, values, attention_mask=None, attention_weights=None)

Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s, nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask: Mask over attention values (b_s, h, nq, nk).

True indicates masking. :param attention_weights: Multiplicative weights for attention values (b_s, h, nq, nk). :return:

Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
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
def forward(self, queries, keys, values, attention_mask=None, attention_weights=None):
    """Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
    nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask:
    Mask over attention values (b_s, h, nq, nk).

    True indicates masking.
    :param attention_weights: Multiplicative weights for attention values (b_s, h,
        nq, nk).
    :return:
    """
    b_s, nq = queries.shape[:2]
    nk = keys.shape[1]

    q = self.fc_q(queries).view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)
    k = self.fc_k(keys).view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)
    v = self.fc_v(values).view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)

    att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, nk)
    if attention_weights is not None:
        att = att * attention_weights
    if attention_mask is not None:
        att = att.masked_fill(attention_mask, -np.inf)
    att = torch.softmax(att, -1)
    att = self.dropout(att)

    out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)
    out = self.fc_o(out)  # (b_s, nq, d_model)
    return out

SimplifiedScaledDotProductAttention

Bases: Module

Scaled dot-product attention.

Methods:

Name Description
__init__

:param d_model: Output dimensionality of the model
forward

Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
class SimplifiedScaledDotProductAttention(nn.Module):
    """Scaled dot-product attention."""

    def __init__(self, d_model, h, dropout=0.1):
        """
        :param d_model: Output dimensionality of the model
        :param d_k: Dimensionality of queries and keys
        :param d_v: Dimensionality of values
        :param h: Number of heads
        """
        super().__init__()

        self.d_model = d_model
        self.d_k = d_model // h
        self.d_v = d_model // h
        self.h = h

        self.fc_o = nn.Linear(h * self.d_v, d_model)
        self.dropout = nn.Dropout(dropout)

        self.init_weights()

    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode="fan_out")
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, queries, keys, values, attention_mask=None, attention_weights=None):
        """Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
        nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask:
        Mask over attention values (b_s, h, nq, nk).

        True indicates masking.
        :param attention_weights: Multiplicative weights for attention values (b_s, h,
            nq, nk).
        :return:
        """
        b_s, nq = queries.shape[:2]
        nk = keys.shape[1]

        q = queries.view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)
        k = keys.view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)
        v = values.view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)

        att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, nk)
        if attention_weights is not None:
            att = att * attention_weights
        if attention_mask is not None:
            att = att.masked_fill(attention_mask, -np.inf)
        att = torch.softmax(att, -1)
        att = self.dropout(att)

        out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)
        out = self.fc_o(out)  # (b_s, nq, d_model)
        return out

__init__ 

__init__(d_model, h, dropout=0.1)

:param d_model: Output dimensionality of the model :param d_k: Dimensionality of queries and keys :param d_v: Dimensionality of values :param h: Number of heads

Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
def __init__(self, d_model, h, dropout=0.1):
    """
    :param d_model: Output dimensionality of the model
    :param d_k: Dimensionality of queries and keys
    :param d_v: Dimensionality of values
    :param h: Number of heads
    """
    super().__init__()

    self.d_model = d_model
    self.d_k = d_model // h
    self.d_v = d_model // h
    self.h = h

    self.fc_o = nn.Linear(h * self.d_v, d_model)
    self.dropout = nn.Dropout(dropout)

    self.init_weights()

forward 

forward(queries, keys, values, attention_mask=None, attention_weights=None)

Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s, nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask: Mask over attention values (b_s, h, nq, nk).

True indicates masking. :param attention_weights: Multiplicative weights for attention values (b_s, h, nq, nk). :return:

Source code in deeplabcut/pose_estimation_pytorch/models/modules/coam_module.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
def forward(self, queries, keys, values, attention_mask=None, attention_weights=None):
    """Computes :param queries: Queries (b_s, nq, d_model) :param keys: Keys (b_s,
    nk, d_model) :param values: Values (b_s, nk, d_model) :param attention_mask:
    Mask over attention values (b_s, h, nq, nk).

    True indicates masking.
    :param attention_weights: Multiplicative weights for attention values (b_s, h,
        nq, nk).
    :return:
    """
    b_s, nq = queries.shape[:2]
    nk = keys.shape[1]

    q = queries.view(b_s, nq, self.h, self.d_k).permute(0, 2, 1, 3)  # (b_s, h, nq, d_k)
    k = keys.view(b_s, nk, self.h, self.d_k).permute(0, 2, 3, 1)  # (b_s, h, d_k, nk)
    v = values.view(b_s, nk, self.h, self.d_v).permute(0, 2, 1, 3)  # (b_s, h, nk, d_v)

    att = torch.matmul(q, k) / np.sqrt(self.d_k)  # (b_s, h, nq, nk)
    if attention_weights is not None:
        att = att * attention_weights
    if attention_mask is not None:
        att = att.masked_fill(attention_mask, -np.inf)
    att = torch.softmax(att, -1)
    att = self.dropout(att)

    out = torch.matmul(att, v).permute(0, 2, 1, 3).contiguous().view(b_s, nq, self.h * self.d_v)  # (b_s, nq, h*d_v)
    out = self.fc_o(out)  # (b_s, nq, d_model)
    return out