Add causal person fusion and MuVAP stage-3 baselines

2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions
@@ -0,0 +1,4 @@
+from .mamba import OADMamba
+from .model_builder import build_model
+from .rnn import MROAD
+from .transformer_models import ViTEnc
@@ -0,0 +1 @@
+from .mamba import OADMamba, OADMambaA
@@ -0,0 +1,172 @@
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from mamba_ssm import Mamba
+# from mambapy.mamba import Mamba, MambaConfig
+from model.model_builder import META_ARCHITECTURES
+from utils.util import parse_flow_type, parse_rgb_type
+
+
+@META_ARCHITECTURES.register("Mamba")
+class OADMamba(nn.Module):
+    
+    def __init__(self, cfg):
+        super(OADMamba, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.input_dim = 0
+        
+        rgb_type = parse_rgb_type(cfg['rgb_type'])
+        flow_type = parse_flow_type(cfg['flow_type'])
+        
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[rgb_type]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[flow_type]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        
+        # raise NotImplementedError("The following code is not implemented yet")
+        self.out_dim = cfg['num_classes']
+        self.window_size = cfg['window_size']
+        
+        self.relu = nn.ReLU()
+        
+        self.embedding_dim = cfg['embedding_dim']
+        self.d_state = cfg['d_state']
+        self.d_conv = cfg['d_conv']
+        self.expand = cfg['expand']
+
+        # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers)
+        # self.mamba = Mamba(self.mamba_config)
+        self.mamba = Mamba(
+            d_model=self.embedding_dim, # Model dimension d_model
+            d_state=self.d_state,  # SSM state expansion factor
+            d_conv=self.d_conv,    # Local convolution width
+            expand=self.expand,    # Block expansion factor
+        )
+        
+        # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout']),
+        )
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim)
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            x = torch.cat((x, flow_input), dim=2)
+        
+        x = self.layer1(x)
+        ht = self.mamba(x)
+        
+        logits = self.f_classification(ht)
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            out_dict['logits'] = pred_scores
+        return out_dict
+
+@META_ARCHITECTURES.register("MambaA")
+class OADMambaA(nn.Module):
+    
+    def __init__(self, cfg):
+        super(OADMambaA, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.input_dim = 0
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[cfg['flow_type']]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.anticipation_length = cfg["anticipation_length"]
+        self.out_dim = cfg['num_classes']
+        
+        self.embedding_dim = cfg['embedding_dim']
+        self.d_state = cfg['d_state'] if 'd_state' in cfg else 16
+        self.d_conv = cfg['d_conv'] if 'd_conv' in cfg else 4
+        self.expand = cfg['expand'] if 'expand' in cfg else 2
+        
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout'])
+        )
+        self.actionness = cfg['actionness']
+        if self.actionness:
+            self.f_actionness = nn.Sequential(
+                nn.Linear(self.hidden_dim, 1),
+            )
+        self.relu = nn.ReLU()
+        # self.mamba_config = MambaConfig(d_model=self.embedding_dim, n_layers=self.num_layers)
+        # self.mamba = Mamba(self.mamba_config)
+        self.mamba = Mamba(
+            d_model=self.embedding_dim, # Model dimension d_model
+            d_state=self.d_state,  # SSM state expansion factor
+            d_conv=self.d_conv,    # Local convolution width
+            expand=self.expand,    # Block expansion facto
+        )
+        
+        # self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.embedding_dim, self.out_dim)
+        )
+        self.anticipation_layer = nn.Sequential(
+            nn.Linear(self.embedding_dim, self.anticipation_length*self.embedding_dim),
+        )
+        self.cfg = cfg
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            # [64, 128, 2048]
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            # [64, 128, 5120]
+            x = torch.cat((x, flow_input), dim=2)
+
+        # audio channel is too huge to directly concat...
+        # to c
+
+        # B, L, D: mamba
+        B, S, feat_dim = x.shape # batch, length, dim
+        # 
+        x = self.layer1(x)
+        ht = self.mamba(x)
+        # h0 = torch.zeros(1, B, self.hidden_dim).to(x.device)
+        # ht, _ = self.gru(x, h0)
+        logits = self.f_classification(self.relu(ht))
+        anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.embedding_dim)
+        anticipation_logits = self.f_classification(self.relu(anticipation_ht))
+        
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+            out_dict['anticipation_logits'] = anticipation_logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1)
+            out_dict['logits'] = pred_scores
+            out_dict['anticipation_logits'] = pred_anticipation_scores
+
+        return out_dict
@@ -0,0 +1,9 @@
+__all__ = ['build_model']
+
+from utils import Registry
+
+META_ARCHITECTURES = Registry()
+
+def build_model(cfg, device=None):
+    model = META_ARCHITECTURES[cfg["model"]](cfg)
+    return model.to(device)
@@ -0,0 +1 @@
+from .rnn import MROAD, MROADA
@@ -0,0 +1,153 @@
+import sys
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from model.model_builder import META_ARCHITECTURES
+from utils.util import parse_flow_type, parse_rgb_type
+
+
+@META_ARCHITECTURES.register("MiniROAD")
+class MROAD(nn.Module):
+    
+    def __init__(self, cfg):
+        super(MROAD, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.input_dim = 0
+        
+        rgb_type = parse_rgb_type(cfg['rgb_type'])
+        flow_type = parse_flow_type(cfg['flow_type'])
+        
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[rgb_type]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[flow_type]
+
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.out_dim = cfg['num_classes']
+        self.window_size = cfg['window_size']
+
+        self.relu = nn.ReLU()
+        self.embedding_dim = cfg['embedding_dim']
+        self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout']),
+        )
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        # self.h0 = torch.nn.Parameter(torch.zeros(1, 1, self.hidden_dim))
+        self.h0 = torch.zeros(self.num_layers, 1, self.hidden_dim)
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            x = torch.cat((x, flow_input), dim=2)
+
+        # if self.use_rgb and self.use_flow:
+        #     x = torch.cat((rgb_input, flow_input), 2)
+        # elif self.use_rgb:
+        #     x = rgb_input
+        # elif self.use_flow:
+        #     x = flow_input
+        
+        x = self.layer1(x)
+        B, _, _ = x.shape
+        h0 = self.h0.expand(-1, B, -1).to(x.device)
+        ht, _ = self.gru(x, h0)
+        ht = self.relu(ht)
+        # ht = self.relu(ht + x)
+        logits = self.f_classification(ht)
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            out_dict['logits'] = pred_scores
+        return out_dict
+
+@META_ARCHITECTURES.register("MiniROADA")
+class MROADA(nn.Module):
+    
+    def __init__(self, cfg):
+        super(MROADA, self).__init__()
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.input_dim = 0
+        if self.use_rgb:
+            self.input_dim += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.input_dim += FEATURE_SIZES[cfg['flow_type']]
+
+        self.embedding_dim = cfg['embedding_dim']
+        self.hidden_dim = cfg['hidden_dim']
+        self.num_layers = cfg['num_layers']
+        self.anticipation_length = cfg["anticipation_length"]
+        self.out_dim = cfg['num_classes']
+
+        self.layer1 = nn.Sequential(
+            nn.Linear(self.input_dim, self.embedding_dim),
+            nn.LayerNorm(self.embedding_dim),
+            nn.ReLU(),
+            nn.Dropout(p=cfg['dropout'])
+        )
+        
+        self.actionness = cfg['actionness']
+        if self.actionness:
+            self.f_actionness = nn.Sequential(
+                nn.Linear(self.hidden_dim, 1),
+            )
+        self.relu = nn.ReLU()
+        self.gru = nn.GRU(self.embedding_dim, self.hidden_dim, self.num_layers, batch_first=True)
+        self.f_classification = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.out_dim)
+        )
+        self.anticipation_layer = nn.Sequential(
+            nn.Linear(self.hidden_dim, self.anticipation_length*self.hidden_dim),
+        )
+        self.cfg = cfg
+
+    def forward(self, rgb_input, flow_input):
+        x = torch.empty(0, device='cuda:0')
+        if self.use_rgb:
+            # [64, 128, 2048]
+            x = torch.cat((x, rgb_input), dim=2)  # Concatenating along the channel dimension
+        if self.use_flow:
+            # [64, 128, 5120]
+            x = torch.cat((x, flow_input), dim=2)
+
+        # audio channel is too huge to directly concat...
+        # to c
+
+        B, S, feat_dim = x.shape # batch, length, dim
+        # 
+        x = self.layer1(x)
+        h0 = torch.zeros(1, B, self.hidden_dim).to(x.device)
+        ht, _ = self.gru(x, h0)
+        ht = self.relu(ht)
+        
+        logits = self.f_classification(self.relu(ht))
+        anticipation_ht = self.anticipation_layer(self.relu(ht)).view(B, S, self.anticipation_length, self.hidden_dim)
+        anticipation_logits = self.f_classification(self.relu(anticipation_ht))
+        
+        
+        out_dict = {}
+        if self.training:
+            out_dict['logits'] = logits
+            out_dict['anticipation_logits'] = anticipation_logits
+        else:
+            pred_scores = F.softmax(logits, dim=-1)
+            pred_anticipation_scores = F.softmax(anticipation_logits, dim=-1)
+            out_dict['logits'] = pred_scores
+            out_dict['anticipation_logits'] = pred_anticipation_scores
+
+        return out_dict
@@ -0,0 +1,155 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SelfAttention(nn.Module):
+    def __init__(
+        self, dim, heads=8, qkv_bias=False, qk_scale=None, dropout_rate=0.0
+    ):
+        super().__init__()
+        self.num_heads = heads
+        head_dim = dim // heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(dropout_rate)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class AxialAttention(nn.Module):
+    def __init__(
+        self,
+        in_planes,
+        out_planes,
+        groups=8,
+        kernel_size=56,
+        stride=1,
+        bias=False,
+        width=False,
+    ):
+        assert (in_planes % groups == 0) and (out_planes % groups == 0)
+        super(AxialAttention, self).__init__()
+        self.in_planes = in_planes
+        self.out_planes = out_planes
+        self.groups = groups
+        self.group_planes = out_planes // groups
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.bias = bias
+        self.width = width
+
+        # Multi-head self attention
+        self.qkv_transform = nn.Conv1d(
+            in_planes,
+            out_planes * 2,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            bias=False,
+        )
+        self.bn_qkv = nn.BatchNorm1d(out_planes * 2)
+        self.bn_similarity = nn.BatchNorm2d(groups * 3)
+        self.bn_output = nn.BatchNorm1d(out_planes * 2)
+
+        # Position embedding
+        self.relative = nn.Parameter(
+            torch.randn(self.group_planes * 2, kernel_size * 2 - 1),
+            requires_grad=True,
+        )
+        query_index = torch.arange(kernel_size).unsqueeze(0)
+        key_index = torch.arange(kernel_size).unsqueeze(1)
+        relative_index = key_index - query_index + kernel_size - 1
+        self.register_buffer('flatten_index', relative_index.view(-1))
+        if stride > 1:
+            self.pooling = nn.AvgPool2d(stride, stride=stride)
+
+        self.reset_parameters()
+
+    def forward(self, x):
+        if self.width:
+            x = x.permute(0, 2, 1, 3)
+        else:
+            x = x.permute(0, 3, 1, 2)  # N, W, C, H
+        N, W, C, H = x.shape
+        x = x.contiguous().view(N * W, C, H)
+
+        # Transformations
+        qkv = self.bn_qkv(self.qkv_transform(x))
+        q, k, v = torch.split(
+            qkv.reshape(N * W, self.groups, self.group_planes * 2, H),
+            [self.group_planes // 2, self.group_planes // 2, self.group_planes],
+            dim=2,
+        )
+
+        # Calculate position embedding
+        all_embeddings = torch.index_select(
+            self.relative, 1, self.flatten_index
+        ).view(self.group_planes * 2, self.kernel_size, self.kernel_size)
+        q_embedding, k_embedding, v_embedding = torch.split(
+            all_embeddings,
+            [self.group_planes // 2, self.group_planes // 2, self.group_planes],
+            dim=0,
+        )
+        qr = torch.einsum('bgci,cij->bgij', q, q_embedding)
+        kr = torch.einsum('bgci,cij->bgij', k, k_embedding).transpose(2, 3)
+        qk = torch.einsum('bgci, bgcj->bgij', q, k)
+        stacked_similarity = torch.cat([qk, qr, kr], dim=1)
+        stacked_similarity = (
+            self.bn_similarity(stacked_similarity)
+            .view(N * W, 3, self.groups, H, H)
+            .sum(dim=1)
+        )
+
+        # (N, groups, H, H, W)
+        similarity = F.softmax(stacked_similarity, dim=3)
+        sv = torch.einsum('bgij,bgcj->bgci', similarity, v)
+        sve = torch.einsum('bgij,cij->bgci', similarity, v_embedding)
+        stacked_output = torch.cat([sv, sve], dim=-1).view(
+            N * W, self.out_planes * 2, H
+        )
+        output = (
+            self.bn_output(stacked_output)
+            .view(N, W, self.out_planes, 2, H)
+            .sum(dim=-2)
+        )
+
+        if self.width:
+            output = output.permute(0, 2, 1, 3)
+        else:
+            output = output.permute(0, 2, 3, 1)
+
+        if self.stride > 1:
+            output = self.pooling(output)
+
+        return output
+
+    def reset_parameters(self):
+        self.qkv_transform.weight.data.normal_(
+            0, math.sqrt(1.0 / self.in_planes)
+        )
+        nn.init.normal_(self.relative, 0.0, math.sqrt(1.0 / self.group_planes))
@@ -0,0 +1,244 @@
+import torch
+import torch.nn as nn
+from .Attention import AxialAttention
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class AxialBlock(nn.Module):
+    expansion = 2
+
+    def __init__(
+        self,
+        inplanes,
+        planes,
+        stride=1,
+        downsample=None,
+        groups=1,
+        base_width=64,
+        dilation=1,
+        norm_layer=None,
+        kernel_size=56,
+    ):
+        super(AxialBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.0))
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv_down = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.hight_block = AxialAttention(
+            width, width, groups=groups, kernel_size=kernel_size
+        )
+        self.width_block = AxialAttention(
+            width,
+            width,
+            groups=groups,
+            kernel_size=kernel_size,
+            stride=stride,
+            width=True,
+        )
+        self.conv_up = conv1x1(width, planes * self.expansion)
+        self.bn2 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv_down(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.hight_block(out)
+        out = self.width_block(out)
+        out = self.relu(out)
+
+        out = self.conv_up(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class AxialAttentionNet(nn.Module):
+    def __init__(
+        self,
+        layers,
+        num_classes=1000,
+        zero_init_residual=True,
+        groups=8,
+        width_per_group=64,
+        replace_stride_with_dilation=None,
+        norm_layer=None,
+        s=0.5,
+    ):
+        super(AxialAttentionNet, self).__init__()
+        block = AxialBlock
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = int(64 * s)
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError(
+                "replace_stride_with_dilation should be None "
+                "or a 3-element tuple, got {}".format(
+                    replace_stride_with_dilation
+                )
+            )
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(
+            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(
+            block, int(128 * s), layers[0], kernel_size=56
+        )
+        self.layer2 = self._make_layer(
+            block,
+            int(256 * s),
+            layers[1],
+            stride=2,
+            kernel_size=56,
+            dilate=replace_stride_with_dilation[0],
+        )
+        self.layer3 = self._make_layer(
+            block,
+            int(512 * s),
+            layers[2],
+            stride=2,
+            kernel_size=28,
+            dilate=replace_stride_with_dilation[1],
+        )
+        self.layer4 = self._make_layer(
+            block,
+            int(1024 * s),
+            layers[3],
+            stride=2,
+            kernel_size=14,
+            dilate=replace_stride_with_dilation[2],
+        )
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(int(1024 * block.expansion * s), num_classes)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+                if isinstance(m, nn.Conv1d):
+                    pass
+                else:
+                    nn.init.kaiming_normal_(
+                        m.weight, mode='fan_out', nonlinearity='relu'
+                    )
+            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, AxialBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(
+        self, block, planes, blocks, kernel_size=56, stride=1, dilate=False
+    ):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(
+                self.inplanes,
+                planes,
+                stride,
+                downsample,
+                groups=self.groups,
+                base_width=self.base_width,
+                dilation=previous_dilation,
+                norm_layer=norm_layer,
+                kernel_size=kernel_size,
+            )
+        )
+        self.inplanes = planes * block.expansion
+        if stride != 1:
+            kernel_size = kernel_size // 2
+
+        for _ in range(1, blocks):
+            layers.append(
+                block(
+                    self.inplanes,
+                    planes,
+                    groups=self.groups,
+                    base_width=self.base_width,
+                    dilation=self.dilation,
+                    norm_layer=norm_layer,
+                    kernel_size=kernel_size,
+                )
+            )
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x, include_conv5=False, include_top=False):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        if include_conv5:
+            x = self.layer4(x)
+
+        if include_top:
+            x = self.avgpool(x)
+            x = torch.flatten(x, 1)
+            x = self.fc(x)
+
+        return x
+
+    def forward(self, x, include_conv5=False, include_top=False):
+        return self._forward_impl(x, include_conv5, include_top)
+
+
+if __name__ == "__main__":
+    model = AxialAttentionNet([1, 2, 4, 1], s=0.5)
+    # model = AxialAttentionNet([3, 4, 6, 3], s=0.25)
+    # print(model)
+
+    x = torch.randn([8, 3, 224, 224])
+    print(model(x).shape)
+    print(model(x, include_conv5=True).shape)
+    print(model(x, include_conv5=True, include_top=True).shape)
@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+
+class StdConv2d(nn.Conv2d):
+    def forward(self, x):
+        w = self.weight
+        v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False)
+        w = (w - m) / torch.sqrt(v + 1e-10)
+        return F.conv2d(
+            x,
+            w,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return StdConv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=dilation,
+        groups=groups,
+        bias=False,
+        dilation=dilation,
+    )
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return StdConv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, bias=False
+    )
+
+
+class PreActBottleneck(nn.Module):
+    """Pre-activation (v2) bottleneck block.
+    Follows the implementation of "Identity Mappings in Deep Residual Networks":
+    https://github.com/KaimingHe/resnet-1k-layers/blob/master/resnet-pre-act.lua
+    Except it puts the stride on 3x3 conv when available.
+    """
+
+    def __init__(self, in_planes, out_planes=None, mid_planes=None, stride=1):
+        super(PreActBottleneck, self).__init__()
+        out_planes = out_planes or in_planes
+        mid_planes = mid_planes or out_planes // 4
+
+        self.gn1 = nn.GroupNorm(32, in_planes)
+        self.conv1 = conv1x1(in_planes, mid_planes)
+        self.gn2 = nn.GroupNorm(32, mid_planes)
+        self.conv2 = conv3x3(mid_planes, mid_planes, stride)
+        self.gn3 = nn.GroupNorm(32, mid_planes)
+        self.conv3 = conv1x1(mid_planes, out_planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride != 1 or in_planes != out_planes:
+            # Projection also with pre-activation according to paper.
+            self.downsample = conv1x1(in_planes, out_planes, stride)
+
+    def forward(self, x):
+        out = self.relu(self.gn1(x))
+
+        # Residual branch
+        residual = x
+        if hasattr(self, 'downsample'):
+            residual = self.downsample(out)
+
+        # Unit's branch
+        out = self.conv1(out)
+        out = self.conv2(self.relu(self.gn2(out)))
+        out = self.conv3(self.relu(self.gn3(out)))
+
+        return out + residual
+
+
+class ResNetV2Model(nn.Module):
+    """Implementation of Pre-activation (v2) ResNet mode."""
+
+    def __init__(self, block_units, width_factor, head_size=21843):
+        super(ResNetV2Model, self).__init__()
+        wf = width_factor  # shortcut 'cause we'll use it a lot.
+
+        # The following will be unreadable if we split lines.
+        # pylint: disable=line-too-long
+        self.conv1 = nn.Sequential(OrderedDict([
+            ('conv', StdConv2d(3, 64*wf, kernel_size=7, stride=2, padding=3, bias=False)),
+            ('pad', nn.ConstantPad2d(1, 0)),
+            ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=0)),
+        ]))
+
+        self.conv2 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=64*wf, out_planes=256*wf, mid_planes=64*wf))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=256*wf, out_planes=256*wf, mid_planes=64*wf)) for i in range(2, block_units[0] + 1)],
+        ))
+        self.conv3 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=256*wf, out_planes=512*wf, mid_planes=128*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=512*wf, out_planes=512*wf, mid_planes=128*wf)) for i in range(2, block_units[1] + 1)],
+        ))
+        self.conv4 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=512*wf, out_planes=1024*wf, mid_planes=256*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=1024*wf, out_planes=1024*wf, mid_planes=256*wf)) for i in range(2, block_units[2] + 1)],
+        ))
+        self.conv5 = nn.Sequential(OrderedDict(
+            [('unit01', PreActBottleneck(in_planes=1024*wf, out_planes=2048*wf, mid_planes=512*wf, stride=2))] +
+            [(f'unit{i:02d}', PreActBottleneck(in_planes=2048*wf, out_planes=2048*wf, mid_planes=512*wf)) for i in range(2, block_units[3] + 1)],
+        ))
+        # pylint: enable=line-too-long
+
+        self.head = nn.Sequential(OrderedDict([
+            ('gn', nn.GroupNorm(32, 2048*wf)),
+            ('relu', nn.ReLU(inplace=True)),
+            ('avg', nn.AdaptiveAvgPool2d(output_size=1)),
+            ('conv', nn.Conv2d(2048*wf, head_size, kernel_size=1, bias=True)),
+        ]))
+
+    def forward(self, x, include_conv5=False, include_top=False):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        if include_conv5:
+            x = self.conv5(x)
+        if include_top:
+            x = self.head(x)
+
+        if include_top and include_conv5:
+            assert x.shape[-2:] == (1,1,)
+            return x[..., 0, 0]
+
+        return x
@@ -0,0 +1,223 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .BiT import ResNetV2Model
+from .AxialNet import AxialAttentionNet
+from .Transformer import TransformerModel
+from .PositionalEncoding import (
+    FixedPositionalEncoding,
+    LearnedPositionalEncoding,
+)
+
+
+class HybridVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5,
+        dropout_rate,
+        positional_encoding_type,
+        backbone=None,
+    ):
+        super(HybridVisionTransformer, self).__init__()
+
+        assert embedding_dim % num_heads == 0
+
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.out_dim = out_dim
+        self.num_channels = num_channels
+        self.include_conv5 = include_conv5
+        self.backbone = backbone
+
+        self.backbone_model, self.flatten_dim = self.configure_backbone()
+
+        self.projection_encoding = nn.Linear(self.flatten_dim, embedding_dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
+
+        self.decoder_dim = int(img_dim / 16.0) ** 2
+        if self.include_conv5:
+            self.decoder_dim = int(img_dim / 32.0) ** 2
+
+        self.decoder_dim += 1  # for the cls token
+
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.decoder_dim, self.embedding_dim, self.decoder_dim
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+
+        self.transformer = TransformerModel(
+            embedding_dim, num_layers, num_heads, hidden_dim
+        )
+        self.mlp_head = nn.Sequential(
+            nn.Linear(embedding_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, out_dim),
+        )
+        self.to_cls_token = nn.Identity()
+
+    def forward(self, x):
+        # apply bit backbone
+        x = self.backbone_model(x, include_conv5=self.include_conv5)
+        x = x.view(x.size(0), -1, self.flatten_dim)
+
+        x = self.projection_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = self.position_encoding(x)
+
+        # apply transformer
+        x = self.transformer(x)
+        x = self.to_cls_token(x[:, 0])
+        x = self.mlp_head(x)
+        x = F.log_softmax(x, dim=-1)
+
+        return x
+
+    def configure_backbone(self):
+        raise NotImplementedError("Method to be called in child class!!")
+
+
+class ResNetHybridViT(HybridVisionTransformer):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5=False,
+        dropout_rate=0.1,
+        positional_encoding_type="learned",
+        backbone='r50x1',
+    ):
+        super(ResNetHybridViT, self).__init__(
+            img_dim=img_dim,
+            out_dim=out_dim,
+            num_channels=num_channels,
+            embedding_dim=embedding_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            hidden_dim=hidden_dim,
+            include_conv5=include_conv5,
+            dropout_rate=dropout_rate,
+            positional_encoding_type=positional_encoding_type,
+            backbone=backbone,
+        )
+
+    def configure_backbone(self):
+        """
+        Current support offered for all BiT models
+        KNOWN_MODELS in https://github.com/google-research/big_transfer/blob/master/bit_pytorch/models.py
+
+        expects model name of style 'r{depth}x{width}'
+        where depth in [50, 101, 152]
+        where width in [1,2,3,4]
+        """
+        backbone = self.backbone
+        out_dim = self.out_dim
+
+        splits = backbone.split('x')
+        model_name = splits[0]
+        width_factor = int(splits[1])
+
+        if model_name in ['r50', 'r101'] and width_factor in [2, 4]:
+            return ValueError(
+                "Invalid Configuration of models -- expect 50x1, 50x3, 101x1, 101x3"
+            )
+        elif model_name == 'r152' and width_factor in [1, 3]:
+            return ValueError(
+                "Invalid Configuration of models -- expect 152x2, 152x4"
+            )
+
+        block_units_dict = {
+            'r50': [3, 4, 6, 3],
+            'r101': [3, 4, 23, 3],
+            'r152': [3, 8, 36, 3],
+        }
+        block_units = block_units_dict.get(model_name, [3, 4, 6, 3])
+        model = ResNetV2Model(block_units, width_factor, head_size=out_dim)
+
+        if self.num_channels == 3:
+            flatten_dim = 1024 * width_factor
+        if self.include_conv5:
+            flatten_dim *= 2
+
+        return model, flatten_dim
+
+
+class AxialNetHybridViT(HybridVisionTransformer):
+    def __init__(
+        self,
+        img_dim,
+        out_dim,
+        num_channels,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        include_conv5=False,
+        dropout_rate=0.1,
+        positional_encoding_type="learned",
+        backbone='a50m',
+    ):
+        super(AxialNetHybridViT, self).__init__(
+            img_dim=img_dim,
+            out_dim=out_dim,
+            num_channels=num_channels,
+            embedding_dim=embedding_dim,
+            num_heads=num_heads,
+            num_layers=num_layers,
+            hidden_dim=hidden_dim,
+            include_conv5=include_conv5,
+            dropout_rate=dropout_rate,
+            positional_encoding_type=positional_encoding_type,
+            backbone=backbone,
+        )
+
+    def configure_backbone(self):
+        """
+        Current support offered for all BiT models
+        models from https://github.com/csrhddlam/axial-deeplab/blob/master/lib/models/axialnet.py
+
+        expects model name of style 'a{depth}{width}'
+        where depth in [26, 50, 101]
+        where width in [s, m, l]
+        """
+        backbone = self.backbone
+        out_dim = self.out_dim
+
+        model_name = backbone[:3]
+        width = backbone[-1]
+
+        block_units_dict = {
+            'a26': [1, 2, 4, 1],
+            'a50': [3, 4, 6, 3],
+            'a101': [3, 4, 23, 3],
+        }
+        block_units = block_units_dict.get(model_name, [3, 4, 6, 3])
+
+        scale_factor_dict = {'s': 0.5, 'm': 0.75, 'l': 1.0}
+        scale_factor = scale_factor_dict.get(width, 0.75)
+        model = AxialAttentionNet(
+            block_units, s=scale_factor, num_classes=out_dim
+        )
+
+        if self.num_channels == 3:
+            flatten_dim = int(512 * float(scale_factor / 0.5))
+        if self.include_conv5:
+            flatten_dim *= 2
+
+        return model, flatten_dim
@@ -0,0 +1,331 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .decoder import Decoder, DecoderLayer
+from .attn import FullAttention, ProbAttention, AttentionLayer
+from .Transformer import TransformerModel
+from ipdb import set_trace
+from .PositionalEncoding import (
+    FixedPositionalEncoding,
+    LearnedPositionalEncoding,
+)
+
+__all__ = ['ViT_B16', 'ViT_B32', 'ViT_L16', 'ViT_L32', 'ViT_H14']
+
+
+class VisionTransformer_v3(nn.Module):
+    def __init__(
+        self,
+        args,
+        img_dim,
+        patch_dim,
+        out_dim,
+        embedding_dim,
+        num_heads,
+        num_layers,
+        hidden_dim,
+        dropout_rate=0.0,
+        attn_dropout_rate=0.0,
+        use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned", with_camera=True, with_motion=True, num_channels=3072,
+    ):
+        super(VisionTransformer_v3, self).__init__()
+
+        assert embedding_dim % num_heads == 0
+        assert img_dim % patch_dim == 0
+        self.with_camera = with_camera
+        self.with_motion = with_motion
+        self.embedding_dim = embedding_dim
+        self.num_heads = num_heads
+        self.patch_dim = patch_dim
+        # num_channels = img_dim
+        self.num_channels = num_channels
+        self.dropout_rate = dropout_rate
+        self.attn_dropout_rate = attn_dropout_rate
+        self.conv_patch_representation = conv_patch_representation
+
+        # self.num_patches = int((img_dim // patch_dim) ** 2)
+        self.num_patches = int(img_dim // patch_dim)
+        self.seq_length = self.num_patches + 1
+        self.flatten_dim = patch_dim * patch_dim * num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            embedding_dim,
+            num_layers,
+            num_heads,
+            hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(embedding_dim)
+
+        d_model = args.decoder_embedding_dim
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(embedding_dim + d_model, hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(hidden_dim//2, out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(embedding_dim + d_model, out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+
+        # Decoder
+        factor = 1  # 5
+        dropout = args.decoder_attn_dropout_rate
+        # d_model = args.decoder_embedding_dim
+        n_heads = args.decoder_num_heads
+        d_layers = args.decoder_layers
+        d_ff = args.decoder_embedding_dim_out  # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None
+        activation = 'gelu'  # 'gelu'
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(FullAttention(True, factor, attention_dropout=dropout),  # True
+                                   d_model, n_heads),  # ProbAttention  FullAttention
+                    AttentionLayer(FullAttention(False, factor, attention_dropout=dropout),  # False
+                                   d_model, n_heads),
+                    d_model,
+                    d_ff,
+                    dropout=dropout,
+                    activation=activation,
+                )
+                for l in range(d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.decoder_cls_token = nn.Parameter(torch.zeros(1, args.query_num, d_model))
+        if positional_encoding_type == "learned":
+            self.decoder_position_encoding = LearnedPositionalEncoding(
+                args.query_num, self.embedding_dim, args.query_num
+            )
+        elif positional_encoding_type == "fixed":
+            self.decoder_position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position decoding :', positional_encoding_type)
+        self.classifier = nn.Linear(d_model, out_dim)
+        self.after_dropout = nn.Dropout(p=self.dropout_rate)
+        # self.merge_fc = nn.Linear(d_model, 1)
+        # self.merge_sigmoid = nn.Sigmoid()
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.with_camera and self.with_motion:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.with_camera:
+            x = sequence_input_rgb
+        elif self.with_motion:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1)
+        x = self.position_encoding(x)
+        x = self.pe_dropout(x)   # not delete
+
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # [128, 33, 1024]
+        # x = self.after_dropout(x)  # add
+        # decoder
+        decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1)
+        # decoder_cls_token = self.after_dropout(decoder_cls_token)  # add
+        # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token)  # [128, 8, 1024]
+        dec = self.decoder(decoder_cls_token, x)   # [128, 8, 1024]
+        dec = self.after_dropout(dec)  # add
+        # merge_atte = self.merge_sigmoid(self.merge_fc(dec))  # [128, 8, 1]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)  # [128, 1024]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001)
+        dec_for_token = dec.mean(dim=1)
+        # dec_for_token = dec.max(dim=1)[0]
+        dec_cls_out = self.classifier(dec)
+        # set_trace()
+        # x = self.to_cls_token(x[:, 0])
+        x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1)
+        x = self.mlp_head(x)
+        # x = F.log_softmax(x, dim=-1)
+
+        return x, dec_cls_out
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)
+
+
+def ViT_B16(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 16
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=768,
+        num_heads=12,
+        num_layers=12,
+        hidden_dim=3072,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_B32(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 32
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=768,
+        num_heads=12,
+        num_layers=12,
+        hidden_dim=3072,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_L16(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 16
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1024,
+        num_heads=16,
+        num_layers=24,
+        hidden_dim=4096,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_L32(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 32
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1024,
+        num_heads=16,
+        num_layers=24,
+        hidden_dim=4096,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
+
+
+def ViT_H14(dataset='imagenet'):
+    if dataset == 'imagenet':
+        img_dim = 224
+        out_dim = 1000
+        patch_dim = 14
+    elif 'cifar' in dataset:
+        img_dim = 32
+        out_dim = 10
+        patch_dim = 4
+
+    return VisionTransformer(
+        img_dim=img_dim,
+        patch_dim=patch_dim,
+        out_dim=out_dim,
+        num_channels=3,
+        embedding_dim=1280,
+        num_heads=16,
+        num_layers=32,
+        hidden_dim=5120,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.0,
+        use_representation=False,
+        conv_patch_representation=False,
+        positional_encoding_type="learned",
+    )
@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+
+
+class FixedPositionalEncoding(nn.Module):
+    def __init__(self, embedding_dim, max_length=5000):
+        super(FixedPositionalEncoding, self).__init__()
+
+        pe = torch.zeros(max_length, embedding_dim)
+        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, embedding_dim, 2).float()
+            * (-torch.log(torch.tensor(10000.0)) / embedding_dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + self.pe[: x.size(0), :]
+        return x
+
+
+class LearnedPositionalEncoding(nn.Module):
+    def __init__(self, max_position_embeddings, embedding_dim, seq_length):
+        super(LearnedPositionalEncoding, self).__init__()
+        self.pe = nn.Embedding(max_position_embeddings, embedding_dim)
+        self.seq_length = seq_length
+
+        self.register_buffer(
+            "position_ids",
+            torch.arange(max_position_embeddings).expand((1, -1)),
+        )
+
+    def forward(self, x, position_ids=None):
+        if position_ids is None:
+            position_ids = self.position_ids[:, : self.seq_length]
+
+        try:
+            position_embeddings = self.pe(position_ids)
+            return x + position_embeddings
+        except:
+            import pdb; pdb.set_trace()
@@ -0,0 +1,82 @@
+from torch import nn
+from .Attention import SelfAttention
+
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(x) + x
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(self.norm(x))
+
+
+class PreNormDrop(nn.Module):
+    def __init__(self, dim, dropout_rate, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.dropout = nn.Dropout(p=dropout_rate)
+        self.fn = fn
+
+    def forward(self, x):
+        return self.dropout(self.fn(self.norm(x)))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout_rate):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(p=dropout_rate),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(p=dropout_rate),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class TransformerModel(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads,
+        mlp_dim,
+        dropout_rate=0.1,
+        attn_dropout_rate=0.1,
+    ):
+        super().__init__()
+        layers = []
+        for _ in range(depth):
+            layers.extend(
+                [
+                    Residual(
+                        PreNormDrop(
+                            dim,
+                            dropout_rate,
+                            SelfAttention(
+                                dim, heads=heads, dropout_rate=attn_dropout_rate
+                            ),
+                        )
+                    ),
+                    Residual(
+                        PreNorm(dim, FeedForward(dim, mlp_dim, dropout_rate))
+                    ),
+                ]
+            )
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.net(x)
@@ -0,0 +1,350 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from configs.feature_sizes import FEATURE_SIZES
+from model.model_builder import META_ARCHITECTURES as registry
+
+from .attn import AttentionLayer, FullAttention, ProbAttention
+from .decoder import Decoder, DecoderLayer
+from .PositionalEncoding import (FixedPositionalEncoding,
+                                 LearnedPositionalEncoding)
+from .Transformer import TransformerModel
+
+
+# @registry.register('TRANSFORMER')
+@registry.register('OADTR')
+class OADTR(nn.Module):
+    def __init__(self, cfg, use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned"
+    ):
+        super(OADTR, self).__init__()
+        
+        self.img_dim = cfg["window_size"] # img_dim
+        self.out_dim = cfg["num_classes"]
+        
+        self.patch_dim = cfg["patch_dim"]
+        self.query_num = cfg["query_num"]
+        
+        self.embedding_dim = cfg["embedding_dim"]
+        self.num_heads = cfg["num_heads"]
+        self.num_layers = cfg["num_layers"]
+        self.hidden_dim = cfg["hidden_dim"]
+        
+        self.dropout_rate = cfg["dropout"]
+        self.attn_dropout_rate = cfg["attn_dropout_rate"]
+        
+        self.conv_patch_representation = conv_patch_representation
+        
+        self.decoder_embedding_dim = cfg["decoder_embedding_dim"]
+        self.decoder_num_heads = cfg["decoder_num_heads"]
+        self.decoder_layers = cfg["decoder_layers"]
+        self.decoder_embedding_dim_out = cfg["decoder_embedding_dim_out"]
+        self.decoder_attn_dropout_rate = cfg["decoder_attn_dropout_rate"]
+        
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        
+        self.num_channels = 0
+        if self.use_rgb:
+            self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.num_channels += FEATURE_SIZES[cfg['flow_type']]
+        
+        assert self.embedding_dim % self.num_heads == 0
+        assert self.img_dim % self.patch_dim == 0
+        
+        self.num_patches = int(self.img_dim // self.patch_dim)
+        self.seq_length = self.num_patches + 1
+        # self.seq_length = self.num_patches + self.img_dim
+        self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
+        # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            self.embedding_dim,
+            self.num_layers,
+            self.num_heads,
+            self.hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
+
+        d_model = self.decoder_embedding_dim
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(self.embedding_dim , self.hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dim//2, self.out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(self.embedding_dim + d_model, self.out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+        
+        # Decoder
+        factor = 1  # 5
+        dropout = self.decoder_attn_dropout_rate
+        # d_model = self.decoder_embedding_dim
+        n_heads = self.decoder_num_heads
+        d_layers = self.decoder_layers
+        d_ff = self.decoder_embedding_dim_out  # args.decoder_embedding_dim_out or 4*args.decoder_embedding_dim None
+        activation = 'gelu'  # 'gelu'
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(FullAttention(True, factor, attention_dropout=dropout),  # True
+                                   d_model, n_heads),  # ProbAttention  FullAttention
+                    AttentionLayer(FullAttention(False, factor, attention_dropout=dropout),  # False
+                                   d_model, n_heads),
+                    d_model,
+                    d_ff,
+                    dropout=dropout,
+                    activation=activation,
+                )
+                for l in range(d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.decoder_cls_token = nn.Parameter(torch.zeros(1, self.query_num, d_model))
+        if positional_encoding_type == "learned":
+            self.decoder_position_encoding = LearnedPositionalEncoding(
+                self.query_num, self.embedding_dim, self.query_num
+            )
+        elif positional_encoding_type == "fixed":
+            self.decoder_position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position decoding :', positional_encoding_type)
+        self.classifier = nn.Linear(d_model, self.out_dim)
+        self.after_dropout = nn.Dropout(p=self.dropout_rate)
+        # self.merge_fc = nn.Linear(d_model, 1)
+        # self.merge_sigmoid = nn.Sigmoid()
+
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.use_rgb and self.use_flow:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.use_rgb:
+            x = sequence_input_rgb
+        elif self.use_flow:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x) # B, 128, 1204
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
+        try:
+            x = self.position_encoding(x) # B, seq+1, 1024
+        except:
+            import pdb; pdb.set_trace()
+        x = self.pe_dropout(x)   # not delete
+        
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # B, seq+1, 1024
+
+        # decoder
+        decoder_cls_token = self.decoder_cls_token.expand(x.shape[0], -1, -1) # [B, 8, 1024]
+        # decoder_cls_token = self.after_dropout(decoder_cls_token)  # add
+        # decoder_cls_token = self.decoder_position_encoding(decoder_cls_token)  # [128, 8, 1024]
+        dec = self.decoder(decoder_cls_token, x)   # [128, 8, 1024]
+        dec = self.after_dropout(dec)  # add
+        
+        # merge_atte = self.merge_sigmoid(self.merge_fc(dec))  # [128, 8, 1]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)  # [128, 1024]
+        # dec_for_token = (merge_atte*dec).sum(dim=1)/(merge_atte.sum(dim=-2) + 0.0001)
+        # *** RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x2048 and 1024x2)
+        dec_for_token = dec.mean(dim=1) # [64, 1024]
+        
+        # dec_for_token = dec.max(dim=1)[0]
+        dec_cls_out = self.classifier(dec) # [B, 8, NUM_CLASS?]
+        # set_trace()
+        # x = self.to_cls_token(x[:, 0])
+        
+        x = torch.cat((self.to_cls_token(x[:, -1]), dec_for_token), dim=1) # [B, 2048]
+        x = self.mlp_head(x) # [B, NUM_CLASSES]
+        
+        # original OADTR
+        # return x, dec_cls_out
+        
+        # x = F.log_softmax(x, dim=-1)
+        out_dict = {}
+        out_dict['logits'] = x.unsqueeze(1) # [B, 1, NUM_CLASSES] ex) [64, 1, 2]
+        out_dict['dec_logits'] = dec_cls_out # [B, 8, NUM_CLASSES] ex) [64, 8, 2]
+        
+        # out_dict['logits'] = x #x.unsqueeze(1)
+        return out_dict
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)
+
+@registry.register('Transformer')
+class ViTEnc(nn.Module):
+    def __init__(self, cfg, use_representation=True,
+        conv_patch_representation=False,
+        positional_encoding_type="learned"
+    ):
+        super(ViTEnc, self).__init__()
+        self.img_dim = cfg["window_size"]
+        self.out_dim = cfg["num_classes"]
+        self.embedding_dim = cfg["embedding_dim"]
+        self.patch_dim = cfg["patch_dim"]
+        self.num_heads = cfg["num_heads"]
+        self.num_layers = cfg["num_layers"]
+        self.hidden_dim = cfg["hidden_dim"]
+        self.dropout_rate = cfg["dropout"]
+        self.use_flow = not cfg['no_flow']
+        self.use_rgb = not cfg['no_rgb']
+        self.num_channels= 0
+        if self.use_rgb:
+            self.num_channels += FEATURE_SIZES[cfg['rgb_type']]
+        if self.use_flow:
+            self.num_channels += FEATURE_SIZES[cfg['flow_type']]
+        self.attn_dropout_rate = cfg["attn_dropout_rate"]
+        assert self.embedding_dim % self.num_heads == 0
+        assert self.img_dim % self.patch_dim == 0
+        self.conv_patch_representation = conv_patch_representation
+        self.num_patches = int(self.img_dim // self.patch_dim)
+        self.seq_length = self.num_patches + 1
+        # self.seq_length = self.num_patches + self.img_dim
+        self.flatten_dim = self.patch_dim * self.patch_dim * self.num_channels
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embedding_dim))
+        # self.cls_token = nn.Parameter(torch.zeros(1, self.img_dim, self.embedding_dim))
+
+        self.linear_encoding = nn.Linear(self.flatten_dim, self.embedding_dim)
+        if positional_encoding_type == "learned":
+            self.position_encoding = LearnedPositionalEncoding(
+                self.seq_length, self.embedding_dim, self.seq_length
+            )
+        elif positional_encoding_type == "fixed":
+            self.position_encoding = FixedPositionalEncoding(
+                self.embedding_dim,
+            )
+        print('position encoding :', positional_encoding_type)
+
+        self.pe_dropout = nn.Dropout(p=self.dropout_rate)
+
+        self.encoder = TransformerModel(
+            self.embedding_dim,
+            self.num_layers,
+            self.num_heads,
+            self.hidden_dim,
+            self.dropout_rate,
+            self.attn_dropout_rate,
+        )
+        self.pre_head_ln = nn.LayerNorm(self.embedding_dim)
+
+        use_representation = False  # False
+        if use_representation:
+            self.mlp_head = nn.Sequential(
+                nn.Linear(self.embedding_dim , self.hidden_dim//2),
+                # nn.Tanh(),
+                nn.ReLU(),
+                nn.Linear(self.hidden_dim//2, self.out_dim),
+            )
+        else:
+            self.mlp_head = nn.Linear(self.embedding_dim, self.out_dim)
+
+        if self.conv_patch_representation:
+            # self.conv_x = nn.Conv2d(
+            #     self.num_channels,
+            #     self.embedding_dim,
+            #     kernel_size=(self.patch_dim, self.patch_dim),
+            #     stride=(self.patch_dim, self.patch_dim),
+            #     padding=self._get_padding(
+            #         'VALID', (self.patch_dim, self.patch_dim),
+            #     ),
+            # )
+            self.conv_x = nn.Conv1d(
+                self.num_channels,
+                self.embedding_dim,
+                kernel_size=self.patch_dim,
+                stride=self.patch_dim,
+                padding=self._get_padding(
+                    'VALID',  (self.patch_dim),
+                ),
+            )
+        else:
+            self.conv_x = None
+
+        self.to_cls_token = nn.Identity()
+
+
+    def forward(self, sequence_input_rgb, sequence_input_flow):
+        if self.use_rgb and self.use_flow:
+            x = torch.cat((sequence_input_rgb, sequence_input_flow), 2)
+        elif self.use_rgb:
+            x = sequence_input_rgb
+        elif self.use_flow:
+            x = sequence_input_flow
+
+        x = self.linear_encoding(x)
+        cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) # B, 1, 1024
+        # x = torch.cat((cls_tokens, x), dim=1)
+        x = torch.cat((x, cls_tokens), dim=1) # B, seq+1, 1024
+        
+        x = self.position_encoding(x) # B, seq+1, 1024
+        x = self.pe_dropout(x)   # not delete
+
+        # apply transformer
+        x = self.encoder(x)
+        x = self.pre_head_ln(x)  # B, seq+1, 1024
+
+        x = self.to_cls_token(x[:, 0]) # B, 1024
+        # x = self.to_cls_token(x[:,0:self.img_dim]) # B, 1024
+        x = self.mlp_head(x)
+        # x = F.log_softmax(x, dim=-1)
+        out_dict = {}
+        out_dict['logits'] = x.unsqueeze(1)
+        # out_dict['logits'] = x #x.unsqueeze(1)
+        return out_dict
+
+    def _get_padding(self, padding_type, kernel_size):
+        assert padding_type in ['SAME', 'VALID']
+        if padding_type == 'SAME':
+            _list = [(k - 1) // 2 for k in kernel_size]
+            return tuple(_list)
+        return tuple(0 for _ in kernel_size)
@@ -0,0 +1,4 @@
+from .ViT import ViTEnc
+from .HybridViT import ResNetHybridViT, AxialNetHybridViT
+
+__all__ = ['ResNetHybridViT', 'AxialNetHybridViT', 'VisionTransformer_v3']
@@ -0,0 +1,170 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+
+from math import sqrt
+
+
+class TriangularCausalMask():
+    def __init__(self, B, L, device="cpu"):
+        mask_shape = [B, 1, L, L]
+        with torch.no_grad():
+            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class ProbMask():
+    def __init__(self, B, H, L, index, scores, device="cpu"):
+        _mask = torch.ones(L, scores.shape[-1], dytpe=torch.bool).to(device).triu(1)
+        _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])
+        indicator = _mask_ex[torch.arange(B)[:, None, None],
+                    torch.arange(H)[None, :, None],
+                    index, :].to(device)
+        self._mask = indicator.view(scores.shape).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class FullAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1):
+        super(FullAttention, self).__init__()
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        scale = self.scale or 1. / sqrt(E)
+
+        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+        if self.mask_flag:
+            if attn_mask is None:
+                attn_mask = TriangularCausalMask(B, L, device=queries.device)
+
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        A = self.dropout(torch.softmax(scale * scores, dim=-1))
+        V = torch.einsum("bhls,bshd->blhd", A, values)
+
+        return V.contiguous()
+
+
+class ProbAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1):
+        super(ProbAttention, self).__init__()
+        self.factor = factor
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def _prob_QK(self, Q, K, sample_k, n_top):
+        # Q [B, H, L, D]
+        B, H, L, E = K.shape
+        _, _, S, _ = Q.shape
+
+        # calculate the sampled Q_K
+        K_expand = K.unsqueeze(-3).expand(B, H, S, L, E)
+        indx_sample = torch.randint(L, (S, sample_k))
+        K_sample = K_expand[:, :, torch.arange(S).unsqueeze(1), indx_sample, :]
+        Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
+
+        # find the Top_k query with sparisty measurement
+        M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L)
+        M_top = M.topk(n_top, sorted=False)[1]
+
+        # use the reduced Q to calculate Q_K
+        Q_reduce = Q[torch.arange(B)[:, None, None],
+                   torch.arange(H)[None, :, None],
+                   M_top, :]
+        Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1))
+
+        return Q_K, M_top
+
+    def _get_initial_context(self, V, L_Q):
+        B, H, L_V, D = V.shape
+        if not self.mask_flag:
+            V_sum = V.sum(dim=-2)
+            contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
+        else:  # use mask
+            assert (L_Q == L_V)  # requires that L_Q == L_V, i.e. for self-attention only
+            contex = V.cumsum(dim=-1)
+        return contex
+
+    def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
+        B, H, L_V, D = V.shape
+
+        if self.mask_flag:
+            attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        attn = torch.softmax(scores, dim=-1)  # nn.Softmax(dim=-1)(scores)
+
+        context_in[torch.arange(B)[:, None, None],
+        torch.arange(H)[None, :, None],
+        index, :] = torch.matmul(attn, V)
+        return context_in
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, D = queries.shape
+        _, S, _, _ = keys.shape
+
+        queries = queries.view(B, H, L, -1)
+        keys = keys.view(B, H, S, -1)
+        values = values.view(B, H, S, -1)
+
+        U = self.factor * np.ceil(np.log(S)).astype('int').item()
+        u = self.factor * np.ceil(np.log(L)).astype('int').item()
+
+        scores_top, index = self._prob_QK(queries, keys, u, U)
+        # add scale factor
+        scale = self.scale or 1. / sqrt(D)
+        if scale is not None:
+            scores_top = scores_top * scale
+        # get the context
+        context = self._get_initial_context(values, L)
+        # update the context with selected top_k queries
+        context = self._update_context(context, values, scores_top, index, L, attn_mask)
+
+        return context.contiguous()
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, attention, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AttentionLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_attention = attention
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out = self.inner_attention(
+            queries,
+            keys,
+            values,
+            attn_mask
+        ).view(B, L, -1)
+
+        return self.out_projection(out)
@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
+                 dropout=0.1, activation="relu"):
+        super(DecoderLayer, self).__init__()
+        d_ff = d_ff or 4*d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask
+        ))
+        x = self.norm1(x)
+
+        x = x + self.dropout(self.cross_attention(
+            x, cross, cross,
+            attn_mask=cross_mask
+        ))
+
+        y = x = self.norm2(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1,1))))
+        y = self.dropout(self.conv2(y).transpose(-1,1))
+
+        return self.norm3(x+y)
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, norm_layer=None):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        for layer in self.layers:
+            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x
@@ -0,0 +1,35 @@
+# FROM LSTR
+import math
+import torch.nn as nn
+
+def weights_init(m):
+    if isinstance(m, nn.Linear):
+        nn.init.kaiming_uniform_(m.weight.data, a=math.sqrt(5))
+    elif isinstance(m, nn.Conv1d):
+        nn.init.normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv2d):
+        nn.init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose1d):
+        nn.init.normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose2d):
+        nn.init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            nn.init.normal_(m.bias.data)
+    elif isinstance(m, nn.BatchNorm1d):
+        nn.init.normal_(m.weight.data, mean=1, std=0.02)
+        nn.init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm2d):
+        nn.init.normal_(m.weight.data, mean=1, std=0.02)
+        nn.init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.GRUCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                nn.init.orthogonal_(param.data)
+            else:
+                nn.init.normal_(param.data)