lib/models/transformers/transformer.py

import copy
from tkinter import N
from typing import Optional, List
import warnings
import torch
import torch.nn.functional as F
from torch import nn, Tensor
import math
from functools import partial
from timm.models.helpers import build_model_with_cfg, named_apply, adapt_input_conv
from timm.models.layers import Mlp, DropPath, trunc_normal_, lecun_normal_
from .position_encoding import build_memory_position_encoding, build_query_position_encoding

from lib.models.transformers.position_encoding import PositionEmbeddingNone,PositionEmbeddingLearned


def with_pos_embed(tensor, pos: Optional[Tensor]):
    return tensor if pos is None else tensor + pos

class MultiheadAttention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.k = nn.Linear(dim, dim, bias=qkv_bias)
        self.v = nn.Linear(dim, dim, bias=qkv_bias)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)


    def forward(self, query, key, value=None,
                attn_mask=None, key_padding_mask=None,
                need_weights=False):
        """
        query: [B, N, C]
        attn_mask: [N, N] torch.float32
        key_padding_mask: [B, N] torch.bool
        """
        B, q_N, C = query.shape
        
        if attn_mask is not None:
            assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
                attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
                'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
            if attn_mask.dtype == torch.uint8:
                warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
                attn_mask = attn_mask.to(torch.bool)

            if attn_mask.dim() == 2:
                attn_mask = attn_mask.unsqueeze(0)
                if list(attn_mask.size()) != [1, query.size(1), key.size(1)]:
                    raise RuntimeError('The size of the 2D attn_mask is not correct.')
            elif attn_mask.dim() == 3:
                if list(attn_mask.size()) != [B * self.num_heads, query.size(1), key.size(1)]:
                    raise RuntimeError('The size of the 3D attn_mask is not correct.')
            else:
                raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
            # attn_mask's dim is 3 now.

        # convert key_padding_mask to bool
        if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
            warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
            key_padding_mask = key_padding_mask.to(torch.bool)

        q = self.q(query).reshape(B, q_N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)  # B,head,HW,C/head
        k = self.k(key).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        v = self.v(value).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)

        attn = (q @ k.transpose(-2, -1)) * self.scale # (B, head, N, N)
        
        if attn_mask is not None:
            if attn_mask.dtype == torch.bool:
                attn.masked_fill_(attn_mask, float('-inf'))
            else:
                attn += attn_mask
        
        if key_padding_mask is not None:
            attn = attn.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),
                float('-inf'),
            )

        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, q_N, C)
        x = self.proj(x)

        if need_weights:
            # average attention weights over heads
            return x, attn.sum(dim=1) / self.num_heads  # attn: (B, q_len, k_len)
        else:
            return x, None


class MLP(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
    """
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): # nn.ReLU
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.drop = nn.Dropout(drop)
        self.fc2 = nn.Linear(hidden_features, out_features)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        return x


class TransformerDecoder(nn.Module):
    def __init__(self, d_model=256, nhead=8, num_decoder_layers=0, mlp_ratio=1, 
                dropout=0., vocab_size=1001, memory_pos_embedding=None, query_pos_embedding=None):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.query_embedding = nn.Embedding(vocab_size, self.d_model)  # (1001, 256)
        self.memory_pos_embedding = memory_pos_embedding
        self.query_pos_embedding = query_pos_embedding

        decoder_layer = TransformerDecoderLayer(d_model=d_model, num_heads=nhead, mlp_ratio=mlp_ratio,
                                                attn_drop=dropout, proj_drop=dropout)
        if num_decoder_layers == 0:
            self.layers = None
        else:
            self.layers = _get_clones(decoder_layer, num_decoder_layers)

        self.num_decoder_layers = num_decoder_layers
        self.scale_factor = float(d_model // nhead) ** 0.5

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out = fan_out // m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()


    def memory_mask_pos_enc(self, attn_mask, feat_sz):
        """
        attn_mask: (B, img_H, img_W)
        feat_sz: feature size
        """
        batch_size = attn_mask.size(0)
        attn_mask = attn_mask.to(torch.float32)
        attn_mask = F.interpolate(attn_mask.unsqueeze(1), size=(feat_sz, feat_sz)).to(torch.bool).squeeze(1)

        pos_embeds = self.memory_pos_embedding(attn_mask) # sine position encoding  (B, C, feat_sz, feat_sz)

        attn_mask = attn_mask.view(batch_size, -1)
        pos_embeds = pos_embeds.view(batch_size, self.d_model, -1).transpose(1, 2)

        return attn_mask, pos_embeds


    def tri_mask(self, length):
        mask = (torch.triu(torch.ones(length, length)) == 1).float().transpose(0, 1)
        mask.masked_fill_(mask == 0, float('-inf'))
        mask.masked_fill_(mask == 1, float(0.))
        return mask


    def forward(self, query, memory, memory_key_padding_mask, memory_pos,
                return_intermediate_output=False, is_inference=False):
        """
        query:        (B, 5, C) or (B, 18, C) or (B, bbox+mask, C)
        memory:       (B, HW, C)
        memory_key_padding_mask:  (B, HW)
        memory_pos:   (B, HW, C)
        """
        # learnable pos encoding
        query_pos = self.query_pos_embedding(query)

        query_mask = self.tri_mask(query.size(1)).to(query.device)

        aux_logits = []
        output = query
        for i in range(len(self.layers)):
            output, attn_weights = self.layers[i](output, memory,
                                                query_pos=query_pos, memory_pos=memory_pos,
                                                query_mask=query_mask,
                                                memory_key_padding_mask=memory_key_padding_mask,
                                                need_weights=False)
            if return_intermediate_output and i < len(self.layers)-1:
                aux_logits.append(output)
        
        if is_inference:
            return output
        else:
            return output, aux_logits


class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model=256, num_heads=8, mlp_ratio=8,
                 attn_drop=0.1, proj_drop=0.1):
        super().__init__()
        # self_attn
        # self.self_attn = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
        # self.dropout1 = nn.Dropout(proj_drop)
        # self.norm1 = nn.LayerNorm(d_model)
        #
        # self.self_attn_det = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
        # self.dropout1_det = nn.Dropout(proj_drop)
        # self.norm1_det = nn.LayerNorm(d_model)

        # cross_attn
        self.cross_attn = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
        self.dropout2 = nn.Dropout(proj_drop)
        self.norm2 = nn.LayerNorm(d_model)

        self.cross_attn_det = MultiheadAttention(dim=d_model, num_heads=num_heads, attn_drop=attn_drop, proj_drop=proj_drop)
        self.dropout2_det = nn.Dropout(proj_drop)
        self.norm2_det = nn.LayerNorm(d_model)

        # mlp
        dim_feedforward = d_model * mlp_ratio
        self.MLP = MLP(in_features=d_model, hidden_features=dim_feedforward, out_features=d_model, drop=proj_drop)
        self.dropout3 = nn.Dropout(proj_drop)
        self.norm3 = nn.LayerNorm(d_model)

        self.MLP_det = MLP(in_features=d_model, hidden_features=dim_feedforward, out_features=d_model, drop=proj_drop)
        self.dropout3_det = nn.Dropout(proj_drop)
        self.norm3_det = nn.LayerNorm(d_model)

        # final
        self.norm_final = nn.LayerNorm(d_model)


    def forward(self, query, memory,
                query_pos: Optional[Tensor] = None,
                memory_pos: Optional[Tensor] = None,
                query_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                need_weights: bool = False):
        """
        query: (B, 5, C)
        memory: (B, HW, C)
        query_pos: (B, 5, C)
        memory_pos: (B, HW, C)
        query_mask: (5,5)
        memory_key_padding_mask: (B, HW)
        """
        # query self attn
        # q = k = with_pos_embed(query, query_pos)
        # query2, selfattn_weights = self.self_attn(q, k, value=query, attn_mask=query_mask, need_weights=need_weights)
        # query = query + self.dropout1(query2)
        # query = self.norm1(query)
        #
        # # memory self attn
        # q = k = with_pos_embed(memory, memory_pos)
        # query2, selfattn_weights = self.self_attn_det(q, k, value=memory, attn_mask=memory_key_padding_mask, need_weights=need_weights)
        # memory = memory + self.dropout1_det(query2)
        # memory = self.norm1_det(memory)

        # query to memory
        query2, crossattn_weights = self.cross_attn(query=with_pos_embed(query, query_pos),
                                                    key=with_pos_embed(memory, memory_pos),
                                                    value=memory,
                                                    key_padding_mask=memory_key_padding_mask,
                                                    need_weights=need_weights)
        query = query + self.dropout2(query2)
        query = self.norm2(query)

        #  memory to query
        query2, crossattn_weights = self.cross_attn_det(query=with_pos_embed(memory, memory_pos),
                                                    key=with_pos_embed(query, query_pos),
                                                    value=query,
                                                    key_padding_mask=memory_key_padding_mask,
                                                    need_weights=need_weights)
        memory = memory + self.dropout2_det(query2)
        memory = self.norm2_det(memory)


        # mlp
        query2 = self.MLP(query)
        query = query + self.dropout3(query2)
        query = self.norm3(query)

        memory = self.MLP_det(memory)
        memory = memory + self.dropout3_det(memory)
        memory = self.norm3_det(memory)

        ### final fusion
        final_x = query[:,4:] + memory[:,4:]
        final = torch.cat([query[:,:4], memory[:,:4] , final_x ],dim=1)
        query = self.norm_final(final)

        if need_weights:
            return query, crossattn_weights
        else:
            return query, None
class TransformerDecoder_v0(nn.Module):
    def __init__(self, d_model=256, nhead=8, num_decoder_layers=0, mlp_ratio=1,
                 dropout=0.):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead

        decoder_layer = TransformerDecoderLayer(d_model=d_model, num_heads=nhead, mlp_ratio=mlp_ratio,
                                                attn_drop=dropout, proj_drop=dropout)
        if num_decoder_layers == 0:
            self.layers = None
        else:
            self.layers = _get_clones(decoder_layer, num_decoder_layers)

        self.num_decoder_layers = num_decoder_layers
        self.scale_factor = float(d_model // nhead) ** 0.5

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out = fan_out // m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()

    def memory_mask_pos_enc(self, attn_mask, feat_sz):
        """
        attn_mask: (B, img_H, img_W)
        feat_sz: feature size
        """
        batch_size = attn_mask.size(0)
        attn_mask = attn_mask.to(torch.float32)
        attn_mask = F.interpolate(attn_mask.unsqueeze(1), size=(feat_sz, feat_sz)).to(torch.bool).squeeze(1)

        pos_embeds = self.memory_pos_embedding(attn_mask)  # sine position encoding  (B, C, feat_sz, feat_sz)

        attn_mask = attn_mask.view(batch_size, -1)
        pos_embeds = pos_embeds.view(batch_size, self.d_model, -1).transpose(1, 2)

        return attn_mask, pos_embeds


    def forward(self, query, memory, query_pos, memory_pos):
        """
        query:        (B, 5, C) or (B, 18, C) or (B, bbox+mask, C)
        memory:       (B, HW, C)
        memory_key_padding_mask:  (B, HW)
        memory_pos:   (B, HW, C)
        """
        # learnable pos encoding
        # query_pos = self.query_pos_embedding(query)

        output = query
        for i in range(len(self.layers)):
            output, attn_weights = self.layers[i](output, memory,
                                                  query_pos=query_pos, memory_pos=memory_pos,
                                                  memory_key_padding_mask=None,
                                                  need_weights=False)


        return output

class VisionLanguageFusionModule(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0., num_vlfusion_layers=0,
                vl_input_type='separate'):
        super().__init__()
        # self.multihead_attn = MultiheadAttention(dim, num_heads=num_heads, 
        #                                     qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop)
        
        self.vl_input_type = vl_input_type
        VLFusion_layer = MultiheadAttention(dim, num_heads=num_heads, 
                                            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=proj_drop)
        if num_vlfusion_layers == 0:
            self.VLFusion_layers = None
        else:
            self.VLFusion_layers = _get_clones(VLFusion_layer, num_vlfusion_layers)

    def forward(self, query, memory,
                query_pos: Optional[Tensor] = None,
                memory_pos: Optional[Tensor] = None,
                query_mask: Optional[Tensor] = None,
                query_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                need_weights: bool = False):

        if self.vl_input_type == 'separate':
            output = query
        elif self.vl_input_type == 'concat':
            output, query_pos = torch.cat([memory, query], dim=1), torch.cat([memory_pos, query_pos], dim=1)
            memory, memory_pos = output.clone(), query_pos.clone()
            memory_key_padding_mask = torch.cat([memory_key_padding_mask, query_key_padding_mask], dim=1)

        for layer in self.VLFusion_layers:
            output, attn_weights = layer(query=with_pos_embed(output, query_pos),
                                        key=with_pos_embed(memory, memory_pos),
                                        value=memory,
                                        key_padding_mask=memory_key_padding_mask,
                                        need_weights=need_weights)      # attn_weights: (B, q_len, k_len)
            
            output = query * output         # todo: 残乘，应该是和前一个block的输出相乘吧，而不是始终和初始的输入相乘？ 或者把这个放在 for循环外边


        return output


def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])


def build_decoder(cfg):
    memory_pos_embedding = build_memory_position_encoding(cfg)
    query_pos_embedding = build_query_position_encoding(cfg)
    return TransformerDecoder(
        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
        dropout=cfg.MODEL.DECODER.DROPOUT,
        nhead=cfg.MODEL.DECODER.NUM_HEADS,
        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
        num_decoder_layers=cfg.MODEL.DECODER.DEC_LAYERS,
        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
        memory_pos_embedding=memory_pos_embedding,
        query_pos_embedding=query_pos_embedding,
    )

from lib.models.transformers.position_encoding import PositionEmbeddingNone,PositionEmbeddingLearned
def build_temporal_decoder(cfg):
    memory_pos_embedding = PositionEmbeddingLearned(15,256)
    query_pos_embedding = PositionEmbeddingLearned(15,256)
    return TransformerDecoder_v0(
        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
        dropout=cfg.MODEL.DECODER.DROPOUT,
        nhead=cfg.MODEL.DECODER.NUM_HEADS,
        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
        num_decoder_layers=2,
        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
        memory_pos_embedding=memory_pos_embedding,
        query_pos_embedding=query_pos_embedding,
    )

def build_text_prompt_decoder(cfg):
    memory_pos_embedding = PositionEmbeddingLearned(15,512)
    query_pos_embedding = PositionEmbeddingLearned(15,512)
    return TransformerDecoder_v0(
        d_model=512,
        dropout=cfg.MODEL.DECODER.DROPOUT,
        nhead=cfg.MODEL.DECODER.NUM_HEADS,
        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
        num_decoder_layers=2,
        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
        memory_pos_embedding=memory_pos_embedding,
        query_pos_embedding=query_pos_embedding,
    )

def build_img_prompt_decoder(cfg):
    memory_pos_embedding = PositionEmbeddingLearned(15,256)
    query_pos_embedding = PositionEmbeddingLearned(15,256)
    return TransformerDecoder_v0(
        d_model=cfg.MODEL.DECODER.HIDDEN_DIM,
        dropout=cfg.MODEL.DECODER.DROPOUT,
        nhead=8,
        mlp_ratio=cfg.MODEL.DECODER.MLP_RATIO,
        num_decoder_layers=2,
        vocab_size=cfg.MODEL.DECODER.VOCAB_SIZE,
        memory_pos_embedding=memory_pos_embedding,
        query_pos_embedding=query_pos_embedding,
    )


def build_rgb_det_decoder(cfg):
    return TransformerDecoder_v0(
        d_model=512,
        dropout=0.1,
        nhead=8,
        mlp_ratio=1,
        num_decoder_layers=1
    )