dlc2action.model.asformer

  1#
  2# Copyright 2020-present by A. Mathis Group and contributors. All rights reserved.
  3#
  4# This project and all its files are licensed under GNU AGPLv3 or later version. A copy is included in dlc2action/LICENSE.AGPL.
  5#
  6# Incorporates code adapted from ASFormer by ChinaYi
  7# Original work Copyright (c) 2021 ChinaYi
  8# Source: https://github.com/ChinaYi/ASFormer
  9# Originally licensed under MIT License
 10# Combined work licensed under GNU AGPLv3
 11#
 12import copy
 13import math
 14from typing import List, Union
 15
 16import numpy as np
 17import torch
 18import torch.nn as nn
 19import torch.nn.functional as F
 20from dlc2action.model.base_model import Model
 21from torch import optim
 22
 23device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 24
 25
 26def exponential_descrease(idx_decoder, p=3):
 27    """Exponential decrease function for the attention window."""
 28    return math.exp(-p * idx_decoder)
 29
 30
 31class AttentionHelper(nn.Module):
 32    def __init__(self):
 33        super(AttentionHelper, self).__init__()
 34        self.softmax = nn.Softmax(dim=-1)
 35
 36    def scalar_dot_att(self, proj_query, proj_key, proj_val, padding_mask):
 37        """
 38        scalar dot attention.
 39        :param proj_query: shape of (B, C, L) => (Batch_Size, Feature_Dimension, Length)
 40        :param proj_key: shape of (B, C, L)
 41        :param proj_val: shape of (B, C, L)
 42        :param padding_mask: shape of (B, C, L)
 43        :return: attention value of shape (B, C, L)
 44        """
 45        m, c1, l1 = proj_query.shape
 46        m, c2, l2 = proj_key.shape
 47
 48        assert c1 == c2
 49
 50        energy = torch.bmm(
 51            proj_query.permute(0, 2, 1), proj_key
 52        )  # out of shape (B, L1, L2)
 53        attention = energy / np.sqrt(c1)
 54        attention = attention + torch.log(
 55            padding_mask + 1e-6
 56        )  # mask the zero paddings. log(1e-6) for zero paddings
 57        attention = self.softmax(attention)
 58        attention = attention * padding_mask
 59        attention = attention.permute(0, 2, 1)
 60        out = torch.bmm(proj_val, attention)
 61        return out, attention
 62
 63
 64class AttLayer(nn.Module):
 65    def __init__(self, q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type):  # r1 = r2
 66        self._fix_types(q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type)
 67        super(AttLayer, self).__init__()
 68
 69        self.query_conv = nn.Conv1d(
 70            in_channels=self.q_dim, out_channels=self.q_dim // self.r1, kernel_size=1
 71        )
 72        self.key_conv = nn.Conv1d(
 73            in_channels=self.k_dim, out_channels=self.k_dim // self.r2, kernel_size=1
 74        )
 75        self.value_conv = nn.Conv1d(
 76            in_channels=self.v_dim, out_channels=self.v_dim // self.r3, kernel_size=1
 77        )
 78
 79        self.conv_out = nn.Conv1d(
 80            in_channels=self.v_dim // self.r3, out_channels=self.v_dim, kernel_size=1
 81        )
 82
 83        assert self.att_type in ["normal_att", "block_att", "sliding_att"]
 84        assert self.stage in ["encoder", "decoder"]
 85
 86        self.att_helper = AttentionHelper()
 87        self.window_mask = self.construct_window_mask()
 88
 89    def _fix_types(self, q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type):
 90        self.q_dim = int(float(q_dim))
 91        self.k_dim = int(float(k_dim))
 92        self.v_dim = int(float(v_dim))
 93        self.r1 = int(float(r1))
 94        self.r2 = int(float(r2))
 95        self.r3 = int(float(r3))
 96        self.bl = int(float(bl))
 97        self.stage = stage
 98        self.att_type = att_type
 99
100    def construct_window_mask(self):
101        """
102        construct window mask of shape (1, l, l + l//2 + l//2), used for sliding window self attention
103        """
104        window_mask = torch.zeros((1, self.bl, self.bl + 2 * (self.bl // 2)))
105        for i in range(self.bl):
106            window_mask[:, :, i : i + self.bl] = 1
107        return window_mask.to(device)
108
109    def forward(self, x1, x2, mask):
110        """Forward pass."""
111        # x1 from the encoder
112        # x2 from the decoder
113
114        query = self.query_conv(x1)
115        key = self.key_conv(x1)
116
117        if self.stage == "decoder":
118            assert x2 is not None
119            value = self.value_conv(x2)
120        else:
121            value = self.value_conv(x1)
122
123        if self.att_type == "normal_att":
124            return self._normal_self_att(query, key, value, mask)
125        elif self.att_type == "block_att":
126            return self._block_wise_self_att(query, key, value, mask)
127        elif self.att_type == "sliding_att":
128            return self._sliding_window_self_att(query, key, value, mask)
129
130    def _normal_self_att(self, q, k, v, mask):
131        m_batchsize, c1, L = q.size()
132        _, c2, L = k.size()
133        _, c3, L = v.size()
134        padding_mask = torch.ones((m_batchsize, 1, L)).to(device) * mask[:, 0:1, :]
135        output, attentions = self.att_helper.scalar_dot_att(q, k, v, padding_mask)
136        output = self.conv_out(F.relu(output))
137        output = output[:, :, 0:L]
138        return output * mask[:, 0:1, :]
139
140    def _block_wise_self_att(self, q, k, v, mask):
141        m_batchsize, c1, L = q.size()
142        _, c2, L = k.size()
143        _, c3, L = v.size()
144
145        nb = L // self.bl
146        if L % self.bl != 0:
147            q = torch.cat(
148                [q, torch.zeros((m_batchsize, c1, self.bl - L % self.bl)).to(device)],
149                dim=-1,
150            )
151            k = torch.cat(
152                [k, torch.zeros((m_batchsize, c2, self.bl - L % self.bl)).to(device)],
153                dim=-1,
154            )
155            v = torch.cat(
156                [v, torch.zeros((m_batchsize, c3, self.bl - L % self.bl)).to(device)],
157                dim=-1,
158            )
159            nb += 1
160
161        padding_mask = torch.cat(
162            [
163                torch.ones((m_batchsize, 1, L)).to(device) * mask[:, 0:1, :],
164                torch.zeros((m_batchsize, 1, self.bl * nb - L)).to(device),
165            ],
166            dim=-1,
167        )
168
169        q = (
170            q.reshape(m_batchsize, c1, nb, self.bl)
171            .permute(0, 2, 1, 3)
172            .reshape(m_batchsize * nb, c1, self.bl)
173        )
174        padding_mask = (
175            padding_mask.reshape(m_batchsize, 1, nb, self.bl)
176            .permute(0, 2, 1, 3)
177            .reshape(m_batchsize * nb, 1, self.bl)
178        )
179        k = (
180            k.reshape(m_batchsize, c2, nb, self.bl)
181            .permute(0, 2, 1, 3)
182            .reshape(m_batchsize * nb, c2, self.bl)
183        )
184        v = (
185            v.reshape(m_batchsize, c3, nb, self.bl)
186            .permute(0, 2, 1, 3)
187            .reshape(m_batchsize * nb, c3, self.bl)
188        )
189
190        output, attentions = self.att_helper.scalar_dot_att(q, k, v, padding_mask)
191        output = self.conv_out(F.relu(output))
192
193        output = (
194            output.reshape(m_batchsize, nb, c3, self.bl)
195            .permute(0, 2, 1, 3)
196            .reshape(m_batchsize, c3, nb * self.bl)
197        )
198        output = output[:, :, 0:L]
199        return output * mask[:, 0:1, :]
200
201    def _sliding_window_self_att(self, q, k, v, mask):
202        m_batchsize, c1, L = q.size()
203        _, c2, _ = k.size()
204        _, c3, _ = v.size()
205
206        # padding zeros for the last segment
207        nb = L // self.bl
208        if L % self.bl != 0:
209            q = torch.cat(
210                [q, torch.zeros((m_batchsize, c1, self.bl - L % self.bl)).to(device)],
211                dim=-1,
212            )
213            k = torch.cat(
214                [k, torch.zeros((m_batchsize, c2, self.bl - L % self.bl)).to(device)],
215                dim=-1,
216            )
217            v = torch.cat(
218                [v, torch.zeros((m_batchsize, c3, self.bl - L % self.bl)).to(device)],
219                dim=-1,
220            )
221            nb += 1
222        padding_mask = torch.cat(
223            [
224                torch.ones((m_batchsize, 1, L)).to(device) * mask[:, 0:1, :],
225                torch.zeros((m_batchsize, 1, self.bl * nb - L)).to(device),
226            ],
227            dim=-1,
228        )
229
230        # sliding window approach, by splitting query_proj and key_proj into shape (c1, l) x (c1, 2l)
231        # sliding window for query_proj: reshape
232        q = (
233            q.reshape(m_batchsize, c1, nb, self.bl)
234            .permute(0, 2, 1, 3)
235            .reshape(m_batchsize * nb, c1, self.bl)
236        )
237
238        # sliding window approach for key_proj
239        # 1. add paddings at the start and end
240        k = torch.cat(
241            [
242                torch.zeros(m_batchsize, c2, self.bl // 2).to(device),
243                k,
244                torch.zeros(m_batchsize, c2, self.bl // 2).to(device),
245            ],
246            dim=-1,
247        )
248        v = torch.cat(
249            [
250                torch.zeros(m_batchsize, c3, self.bl // 2).to(device),
251                v,
252                torch.zeros(m_batchsize, c3, self.bl // 2).to(device),
253            ],
254            dim=-1,
255        )
256        padding_mask = torch.cat(
257            [
258                torch.zeros(m_batchsize, 1, self.bl // 2).to(device),
259                padding_mask,
260                torch.zeros(m_batchsize, 1, self.bl // 2).to(device),
261            ],
262            dim=-1,
263        )
264
265        # 2. reshape key_proj of shape (m_batchsize*nb, c1, 2*self.bl)
266        k = torch.cat(
267            [
268                k[:, :, i * self.bl : (i + 1) * self.bl + (self.bl // 2) * 2]
269                for i in range(nb)
270            ],
271            dim=0,
272        )  # special case when self.bl = 1
273        v = torch.cat(
274            [
275                v[:, :, i * self.bl : (i + 1) * self.bl + (self.bl // 2) * 2]
276                for i in range(nb)
277            ],
278            dim=0,
279        )
280        # 3. construct window mask of shape (1, l, 2l), and use it to generate final mask
281        padding_mask = torch.cat(
282            [
283                padding_mask[:, :, i * self.bl : (i + 1) * self.bl + (self.bl // 2) * 2]
284                for i in range(nb)
285            ],
286            dim=0,
287        )  # of shape (m*nb, 1, 2l)
288        final_mask = self.window_mask.repeat(m_batchsize * nb, 1, 1) * padding_mask
289
290        output, attention = self.att_helper.scalar_dot_att(q, k, v, final_mask)
291        output = self.conv_out(F.relu(output))
292
293        output = (
294            output.reshape(m_batchsize, nb, -1, self.bl)
295            .permute(0, 2, 1, 3)
296            .reshape(m_batchsize, -1, nb * self.bl)
297        )
298        output = output[:, :, 0:L]
299        return output * mask[:, 0:1, :]
300
301
302class MultiHeadAttLayer(nn.Module):
303    def __init__(self, q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type, num_head):
304        super(MultiHeadAttLayer, self).__init__()
305        #         assert v_dim % num_head == 0
306        self.conv_out = nn.Conv1d(v_dim * num_head, v_dim, 1)
307        self.layers = nn.ModuleList(
308            [
309                copy.deepcopy(
310                    AttLayer(q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type)
311                )
312                for i in range(num_head)
313            ]
314        )
315        self.dropout = nn.Dropout(p=0.5)
316
317    def forward(self, x1, x2, mask):
318        """Forward pass."""
319        out = torch.cat([layer(x1, x2, mask) for layer in self.layers], dim=1)
320        out = self.conv_out(self.dropout(out))
321        return out
322
323
324class ConvFeedForward(nn.Module):
325    def __init__(self, dilation, in_channels, out_channels):
326        super(ConvFeedForward, self).__init__()
327        self.layer = nn.Sequential(
328            nn.Conv1d(
329                in_channels, out_channels, 3, padding=dilation, dilation=dilation
330            ),
331            nn.ReLU(),
332        )
333
334    def forward(self, x):
335        """Forward pass."""
336        return self.layer(x)
337
338
339class FCFeedForward(nn.Module):
340    def __init__(self, in_channels, out_channels):
341        super(FCFeedForward, self).__init__()
342        self.layer = nn.Sequential(
343            nn.Conv1d(in_channels, out_channels, 1),  # conv1d equals fc
344            nn.ReLU(),
345            nn.Dropout(),
346            nn.Conv1d(out_channels, out_channels, 1),
347        )
348
349    def forward(self, x):
350        """Forward pass."""
351        return self.layer(x)
352
353
354class AttModule(nn.Module):
355    def __init__(
356        self, dilation, in_channels, out_channels, r1, r2, att_type, stage, alpha
357    ):
358        super(AttModule, self).__init__()
359        self.feed_forward = ConvFeedForward(dilation, in_channels, out_channels)
360        self.instance_norm = nn.InstanceNorm1d(in_channels, track_running_stats=False)
361        self.att_layer = AttLayer(
362            in_channels,
363            in_channels,
364            out_channels,
365            r1,
366            r1,
367            r2,
368            dilation,
369            att_type=att_type,
370            stage=stage,
371        )  # dilation
372        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
373        self.dropout = nn.Dropout()
374        self.alpha = alpha
375
376    def forward(self, x, f, mask):
377        """Forward pass."""
378        out = self.feed_forward(x)
379        out = self.alpha * self.att_layer(self.instance_norm(out), f, mask) + out
380        out = self.conv_1x1(out)
381        out = self.dropout(out)
382        return (x + out) * mask[:, 0:1, :]
383
384
385class PositionalEncoding(nn.Module):
386    "Implement the PE function."
387
388    def __init__(self, d_model, max_len=10000):
389        super(PositionalEncoding, self).__init__()
390        # Compute the positional encodings once in log space.
391        pe = torch.zeros(max_len, d_model)
392        position = torch.arange(0, max_len).unsqueeze(1)
393        div_term = torch.exp(
394            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
395        )
396        pe[:, 0::2] = torch.sin(position * div_term)
397        pe[:, 1::2] = torch.cos(position * div_term)
398        pe = pe.unsqueeze(0).permute(0, 2, 1)  # of shape (1, d_model, l)
399        self.pe = nn.Parameter(pe, requires_grad=True)
400
401    #         self.register_buffer('pe', pe)
402
403    def forward(self, x):
404        """Forward pass."""
405        return x + self.pe[:, :, 0 : x.shape[2]]
406
407
408class Encoder(nn.Module):
409    def __init__(
410        self,
411        num_layers,
412        r1,
413        r2,
414        num_f_maps,
415        input_dim,
416        channel_masking_rate,
417        att_type,
418        alpha,
419    ):
420        super(Encoder, self).__init__()
421        self.conv_1x1 = nn.Conv1d(input_dim, num_f_maps, 1)  # fc layer
422        self.layers = nn.ModuleList(
423            [
424                AttModule(
425                    2**i, num_f_maps, num_f_maps, r1, r2, att_type, "encoder", alpha
426                )
427                for i in range(num_layers)  # 2**i
428            ]
429        )
430
431        self.dropout = nn.Dropout2d(p=channel_masking_rate)
432        self.channel_masking_rate = channel_masking_rate
433
434    def forward(self, x, mask):
435        """Forward pass."""
436        if self.channel_masking_rate > 0:
437            x = x.unsqueeze(2)
438            x = self.dropout(x)
439            x = x.squeeze(2)
440
441        feature = self.conv_1x1(x)
442        for layer in self.layers:
443            feature = layer(feature, None, mask)
444
445        return feature
446
447
448class Decoder(nn.Module):
449    def __init__(
450        self, num_layers, r1, r2, num_f_maps, input_dim, num_classes, att_type, alpha
451    ):
452        super(
453            Decoder, self
454        ).__init__()  # self.position_en = PositionalEncoding(d_model=num_f_maps)
455        self.conv_1x1 = nn.Conv1d(input_dim, num_f_maps, 1)
456        self.layers = nn.ModuleList(
457            [
458                AttModule(
459                    2**i, num_f_maps, num_f_maps, r1, r2, att_type, "decoder", alpha
460                )
461                for i in range(num_layers)  # 2 ** i
462            ]
463        )
464        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
465
466    def forward(self, x, fencoder, mask):
467        """Forward pass."""
468        feature = self.conv_1x1(x)
469        for layer in self.layers:
470            feature = layer(feature, fencoder, mask)
471
472        out = self.conv_out(feature) * mask[:, 0:1, :]
473
474        return out, feature
475
476
477class MyTransformer(nn.Module):
478    def __init__(
479        self,
480        num_layers,
481        r1,
482        r2,
483        num_f_maps,
484        input_dim,
485        channel_masking_rate,
486    ):
487        super(MyTransformer, self).__init__()
488        self.encoder = Encoder(
489            num_layers,
490            r1,
491            r2,
492            num_f_maps,
493            input_dim,
494            channel_masking_rate,
495            att_type="sliding_att",
496            alpha=1,
497        )
498
499    def forward(self, x):
500        """Forward pass."""
501        mask = (x.sum(1).unsqueeze(1) != 0).int()
502        feature = self.encoder(x, mask)
503        feature = feature * mask
504
505        return feature
506
507
508class Predictor(nn.Module):
509    def __init__(
510        self,
511        num_layers,
512        r1,
513        r2,
514        num_f_maps,
515        num_classes,
516        num_decoders,
517    ):
518        super(Predictor, self).__init__()
519        self.decoders = nn.ModuleList(
520            [
521                copy.deepcopy(
522                    Decoder(
523                        num_layers,
524                        r1,
525                        r2,
526                        num_f_maps,
527                        num_classes,
528                        num_classes,
529                        att_type="sliding_att",
530                        alpha=exponential_descrease(s),
531                    )
532                )
533                for s in range(num_decoders)
534            ]
535        )  # num_decoders
536
537        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
538
539    def forward(self, x):
540        """Forward pass."""
541        mask = (x.sum(1).unsqueeze(1) != 0).int()
542        out = self.conv_out(x) * mask[:, 0:1, :]
543        outputs = out.unsqueeze(0)
544
545        for decoder in self.decoders:
546            out, x = decoder(
547                F.softmax(out, dim=1) * mask[:, 0:1, :], x * mask[:, 0:1, :], mask
548            )
549            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
550        return outputs
551
552
553class ASFormer(Model):
554    """
555    An implementation of ASFormer
556    """
557
558    def __init__(
559        self,
560        num_decoders:int,
561        num_layers:int,
562        r1:float,
563        r2:float,
564        num_f_maps:int,
565        input_dim:dict,
566        num_classes:int,
567        channel_masking_rate:float,
568        state_dict_path:str=None,
569        ssl_constructors:List=None,
570        ssl_types:List=None,
571        ssl_modules:List=None,
572    ):
573        input_dim = sum([x[0] for x in input_dim.values()])
574        self.num_f_maps = int(float(num_f_maps))
575        self.params = {
576            "num_layers": int(float(num_layers)),
577            "r1": float(r1),
578            "r2": float(r2),
579            "num_f_maps": self.num_f_maps,
580            "input_dim": int(float(input_dim)),
581            "channel_masking_rate": float(channel_masking_rate),
582        }
583        self.params_predictor = {
584            "num_layers": int(float(num_layers)),
585            "r1": r1,
586            "r2": r2,
587            "num_f_maps": self.num_f_maps,
588            "num_classes": int(float(num_classes)),
589            "num_decoders": int(float(num_decoders)),
590        }
591        super().__init__(ssl_constructors, ssl_modules, ssl_types, state_dict_path)
592
593    def _feature_extractor(self) -> Union[torch.nn.Module, List]:
594        return MyTransformer(**self.params)
595
596    def _predictor(self) -> torch.nn.Module:
597        return Predictor(**self.params_predictor)
598
599    def features_shape(self) -> torch.Size:
600        return torch.Size([self.num_f_maps])