MQ-Det/maskrcnn_benchmark/modeling/rpn/qv_fuse.py

import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from maskrcnn_benchmark.modeling.language_backbone.clip_model import QuickGELU, LayerNorm, DropPath
from maskrcnn_benchmark.modeling.utils import cat, concat_box_prediction_layers, permute_and_flatten
import os

class CrossMultiHeadAttention(nn.Module):
    def __init__(self, v_dim, embed_dim, num_heads, dropout=0.1, cfg=None):
        super(CrossMultiHeadAttention, self).__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.v_dim = v_dim


        assert (
                self.head_dim * self.num_heads == self.embed_dim
        ), f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`: {self.num_heads})."
        self.scale = self.head_dim ** (-0.5)
        self.dropout = dropout

        self.v_proj = nn.Linear(self.v_dim, self.embed_dim)
        self.cache_proj = nn.Linear(self.v_dim, self.embed_dim)
        self.values_cache_proj = nn.Linear(self.v_dim, self.embed_dim)

        self.out_v_proj = nn.Linear(self.embed_dim, self.v_dim)

        self.clamp_min_for_underflow = cfg.MODEL.DYHEAD.FUSE_CONFIG.CLAMP_MIN_FOR_UNDERFLOW
        self.clamp_max_for_overflow = cfg.MODEL.DYHEAD.FUSE_CONFIG.CLAMP_MAX_FOR_OVERFLOW

        self._reset_parameters()

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def _reset_parameters(self):
        nn.init.xavier_uniform_(self.v_proj.weight)
        self.v_proj.bias.data.fill_(0)
        nn.init.xavier_uniform_(self.out_v_proj.weight)
        self.out_v_proj.bias.data.fill_(0)

    def forward(self, v, cache, attention_mask_cache=None):
        bsz, tgt_len, embed_dim = v.size()
        cache = cache.to(v.dtype)

        query_states = self.v_proj(v) * self.scale
        key_states = self._shape(self.cache_proj(cache), -1, bsz)
        value_cache_states = self._shape(self.values_cache_proj(cache), -1, bsz)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states = key_states.view(*proj_shape)
        value_cache_states = value_cache_states.view(*proj_shape)

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {attn_weights.size()}"
            )

        # attn_weights_l = nn.functional.softmax(attn_weights.transpose(1, 2), dim=-1)
        
        if self.clamp_min_for_underflow:
            attn_weights = torch.clamp(attn_weights, min=-50000) # Do not increase -50000, data type half has quite limited range
        if self.clamp_max_for_overflow:
            attn_weights = torch.clamp(attn_weights, max=50000) # Do not increase 50000, data type half has quite limited range

        if attention_mask_cache is not None:
            assert (attention_mask_cache.dim() == 2)
            attention_mask_cache=attention_mask_cache.to(torch.float)
            attention_mask = attention_mask_cache.unsqueeze(1).unsqueeze(1)
            attention_mask = attention_mask.expand(bsz, 1, tgt_len, src_len)
            attention_mask = attention_mask.masked_fill(attention_mask == 0, -9e15)

            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_weights_v = nn.functional.softmax(attn_weights, dim=-1)

        attn_probs_v = F.dropout(attn_weights_v, p=self.dropout, training=self.training)

        attn_output_v = torch.bmm(attn_probs_v, value_cache_states)


        if attn_output_v.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output_v` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is {attn_output_v.size()}"
            )

        attn_output_v = attn_output_v.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output_v = attn_output_v.transpose(1, 2)
        attn_output_v = attn_output_v.reshape(bsz, tgt_len, self.embed_dim)

        attn_output_v = self.out_v_proj(attn_output_v)

        return attn_output_v



class GatedCrossAttentionBlock(nn.Module):
    def __init__(self, v_dim, embed_dim, num_heads, dropout=0.1,
                 drop_path=.0, init_values=1e-4, cfg=None):
        """
        Inputs:
            embed_dim - Dimensionality of input and attention feature vectors
            hidden_dim - Dimensionality of hidden layer in feed-forward network
                         (usually 2-4x larger than embed_dim)
            num_heads - Number of heads to use in the Multi-Head Attention block
            dropout - Amount of dropout to apply in the feed-forward network
        """
        super(GatedCrossAttentionBlock, self).__init__()

        # pre layer norm
        self.layer_norm_v = nn.LayerNorm(v_dim)
        self.layer_norm_c = nn.LayerNorm(v_dim)
        self.attn = CrossMultiHeadAttention(v_dim=v_dim,
                                         embed_dim=embed_dim,
                                         num_heads=num_heads,
                                         dropout=dropout,
                                         cfg=cfg)

        # add layer scale for training stability
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        # self.gamma_v = nn.Parameter(init_values * torch.ones((v_dim)), requires_grad=True)
        if cfg.MODEL.DYHEAD.FUSE_CONFIG.CONDITIONAL_GATE:
            self.conditioned_gamma_v = nn.Sequential(nn.Linear(v_dim, int(v_dim/2)), nn.ReLU(), nn.Linear(int(v_dim/2), v_dim))
            # zero init
            nn.init.xavier_uniform_(self.conditioned_gamma_v[0].weight)
            self.conditioned_gamma_v[0].bias.data.fill_(0)
            self.conditioned_gamma_v[2].weight.data.fill_(0)
            self.conditioned_gamma_v[2].bias.data.fill_(0)
        else:
            self.gamma_v = nn.Parameter(init_values * torch.ones((v_dim)), requires_grad=True)

        self.cfg = cfg

    def forward(self, q0, q1, q2, q3, q4, cache=None, attention_mask_cache=None):
        # aggrage scales
        visu_feat = []
        size_per_level, visual_features_flatten = [], []
        for ii, feat_per_level in enumerate([q0, q1, q2, q3, q4]):
            bs, c, h, w = feat_per_level.shape
            size_per_level.append([h, w])
            feat = permute_and_flatten(feat_per_level, bs, 1, c, h, w)
            visual_features_flatten.append(feat)
        visual_features_flatten = cat(visual_features_flatten, dim=1)


        new_v = self.single_attention_call(visual_features_flatten, cache, attention_mask_cache=attention_mask_cache)

        # [bs, N, C] -> [bs, C, N]
        new_v = new_v.transpose(1, 2).contiguous()

        # recover scales
        start = 0
        for (h, w) in size_per_level:
            new_v_per_level = new_v[:, :, start:start + h * w].view(bs, -1, h, w).contiguous()
            visu_feat.append(new_v_per_level)
            start += h * w

        return visu_feat[0], visu_feat[1], visu_feat[2], visu_feat[3], visu_feat[4]

    def single_attention_call(self, v, cache, attention_mask_cache=None):
        delta_v = self.attn(self.layer_norm_v(v), self.layer_norm_c(cache), attention_mask_cache=attention_mask_cache)
        if self.cfg.MODEL.DYHEAD.FUSE_CONFIG.CONDITIONAL_GATE:
            gate=self.conditioned_gamma_v(v).tanh()
        else:
            gate=self.gamma_v.tanh()
        v = v + self.drop_path(gate * delta_v)
        return v

class QVFuse(torch.nn.Module):
    """
    Early Fusion Module
    """

    def __init__(self, cfg):
        super(QVFuse, self).__init__()
        self.init_configs(cfg)
        self.cfg = cfg

        self.use_checkpoint = False
        if hasattr(cfg.MODEL.DYHEAD, 'USE_CHECKPOINT'):
            self.use_checkpoint = cfg.MODEL.DYHEAD.USE_CHECKPOINT

        # early fusion module
        self.cross_attn = GatedCrossAttentionBlock(v_dim=self.joint_embedding_size,
                    embed_dim=self.embed_dim,
                    num_heads=self.n_head,
                    dropout=0.1,
                    drop_path=.0,
                    init_values=0.,
                    cfg=cfg
                    )
    def init_configs(self, cfg):
        # common params
        self.joint_embedding_size = cfg.MODEL.DYHEAD.FUSE_CONFIG.JOINT_EMB_SIZE
        # mha params
        self.n_head = 4
        self.embed_dim = 512

    def forward(self, x):
        visual_features = x["visual"]
        cache = x["cache"]
        fused_visual_features = None

        if self.use_checkpoint:
            q0, q1, q2, q3, q4 = checkpoint.checkpoint(self.cross_attn,
                visual_features[0], visual_features[1],
                visual_features[2], visual_features[3],
                visual_features[4], cache['cache'], cache['attention_mask']
            )
        else:
            q0, q1, q2, q3, q4 = self.cross_attn(
                visual_features[0], visual_features[1],
                visual_features[2], visual_features[3],
                visual_features[4], cache['cache'], cache['attention_mask']
            )

        fused_visual_features = [q0, q1, q2, q3, q4]

        # return features_dict
        x.update({"visual": fused_visual_features})
        return x