Spaces:
Paused
Paused
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from timm.models.layers import DropPath | |
| class VLFuse(torch.nn.Module): | |
| """ | |
| Early Fusion Module | |
| """ | |
| def __init__(self, ): | |
| super(VLFuse, self).__init__() | |
| self.init_configs() | |
| # early fusion module | |
| # bi-direction (text->image, image->text) | |
| self.b_attn = BiAttentionBlockForCheckpoint(v_dim=self.img_dim, # 256 | |
| l_dim=self.lang_dim, # 768 | |
| embed_dim=self.embed_dim, # 2048 | |
| num_heads=self.n_head, # 8 | |
| dropout=0.1, | |
| drop_path=.0, | |
| init_values=1.0 / 6, | |
| ) | |
| def init_configs(self, ): | |
| # common params | |
| self.img_dim = 256 | |
| self.max_query_len = 256 | |
| self.n_layers =1 | |
| # mha params | |
| self.n_head = 8 | |
| self.embed_dim = 2048 # 2048 by default | |
| self.lang_dim = 256 | |
| def forward(self, x, task=None): | |
| visual_features = x["visual"] | |
| language_dict_features = x["lang"] | |
| fused_visual_features, language_features = self.b_attn( | |
| visual_features, language_dict_features['hidden'], language_dict_features['masks'], task) | |
| language_dict_features['hidden'] = language_features | |
| fused_language_dict_features = language_dict_features | |
| features_dict = {"visual": fused_visual_features, | |
| "lang": fused_language_dict_features} | |
| return features_dict | |
| def masks_to_boxes(masks): | |
| """Compute the bounding boxes around the provided masks | |
| The masks should be in format [N, H, W] where N is the number of masks, (H, W) are the spatial dimensions. | |
| Returns a [N, 4] tensors, with the boxes in xyxy format | |
| """ | |
| if masks.numel() == 0: | |
| return torch.zeros((0, 4), device=masks.device) | |
| h, w = masks.shape[-2:] | |
| y = torch.arange(0, h, dtype=torch.float, device=masks.device) | |
| x = torch.arange(0, w, dtype=torch.float, device=masks.device) | |
| y, x = torch.meshgrid(y, x) | |
| x_mask = (masks * x.unsqueeze(0)) | |
| x_max = x_mask.flatten(1).max(-1)[0] | |
| x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] | |
| y_mask = (masks * y.unsqueeze(0)) | |
| y_max = y_mask.flatten(1).max(-1)[0] | |
| y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0] | |
| return torch.stack([x_min, y_min, x_max, y_max], 1) | |
| class FeatureFuser(nn.Module): | |
| """ | |
| Feature Fuser for SOT (inspired by CondInst) | |
| """ | |
| def __init__(self, in_channels, channels=256): | |
| super().__init__() | |
| self.refine = nn.ModuleList() | |
| for in_channel in in_channels: | |
| self.refine.append(nn.Conv2d(in_channel, channels, 3, padding=1)) | |
| def forward(self, features): | |
| # -4, -3, -2, -1 corresponds to P3, P4, P5, P6 | |
| for i, f in enumerate([-3, -2, -1]): | |
| if i == 0: | |
| x = self.refine[i](features[f]) | |
| else: | |
| x_p = self.refine[i](features[f]) | |
| target_h, target_w = x.size()[2:] | |
| h, w = x_p.size()[2:] | |
| assert target_h % h == 0 | |
| assert target_w % w == 0 | |
| factor_h, factor_w = target_h // h, target_w // w | |
| assert factor_h == factor_w | |
| x_p = aligned_bilinear(x_p, factor_h) | |
| x = x + x_p | |
| return x | |
| def aligned_bilinear(tensor, factor): | |
| assert tensor.dim() == 4 | |
| assert factor >= 1 | |
| assert int(factor) == factor | |
| if factor == 1: | |
| return tensor | |
| h, w = tensor.size()[2:] | |
| tensor = F.pad(tensor, pad=(0, 1, 0, 1), mode="replicate") | |
| oh = factor * h + 1 | |
| ow = factor * w + 1 | |
| tensor = F.interpolate( | |
| tensor, size=(oh, ow), | |
| mode='bilinear', | |
| align_corners=True | |
| ) | |
| tensor = F.pad( | |
| tensor, pad=(factor // 2, 0, factor // 2, 0), | |
| mode="replicate" | |
| ) | |
| return tensor[:, :, :oh - 1, :ow - 1] | |
| class BiMultiHeadAttention(nn.Module): | |
| def __init__(self, v_dim, l_dim, embed_dim, num_heads, dropout=0.1): | |
| super(BiMultiHeadAttention, self).__init__() | |
| self.embed_dim = embed_dim | |
| self.num_heads = num_heads | |
| self.head_dim = embed_dim // num_heads | |
| self.v_dim = v_dim | |
| self.l_dim = l_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.l_proj = nn.Linear(self.l_dim, self.embed_dim) | |
| self.values_v_proj = nn.Linear(self.v_dim, self.embed_dim) | |
| self.values_l_proj = nn.Linear(self.l_dim, self.embed_dim) | |
| self.out_v_proj = nn.Linear(self.embed_dim, self.v_dim) | |
| self.out_l_proj = nn.Linear(self.embed_dim, self.l_dim) | |
| self.stable_softmax_2d = False | |
| self.clamp_min_for_underflow = True | |
| self.clamp_max_for_overflow = True | |
| 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.l_proj.weight) | |
| self.l_proj.bias.data.fill_(0) | |
| nn.init.xavier_uniform_(self.values_v_proj.weight) | |
| self.values_v_proj.bias.data.fill_(0) | |
| nn.init.xavier_uniform_(self.values_l_proj.weight) | |
| self.values_l_proj.bias.data.fill_(0) | |
| nn.init.xavier_uniform_(self.out_v_proj.weight) | |
| self.out_v_proj.bias.data.fill_(0) | |
| nn.init.xavier_uniform_(self.out_l_proj.weight) | |
| self.out_l_proj.bias.data.fill_(0) | |
| def forward(self, v, l, attention_mask_l=None): | |
| bsz, tgt_len, embed_dim = v.size() | |
| query_states = self.v_proj(v) * self.scale | |
| key_states = self._shape(self.l_proj(l), -1, bsz) | |
| value_v_states = self._shape(self.values_v_proj(v), -1, bsz) | |
| value_l_states = self._shape(self.values_l_proj(l), -1, bsz) | |
| proj_shape = (bsz * self.num_heads, -1, self.head_dim) # (bs * 8, -1, embed_dim//8) | |
| query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape) # (bs * 8, seq_len_img, embed_dim//8) | |
| key_states = key_states.view(*proj_shape) # (bs * 8, seq_len_text, embed_dim//8) | |
| value_v_states = value_v_states.view(*proj_shape) | |
| value_l_states = value_l_states.view(*proj_shape) | |
| src_len = key_states.size(1) | |
| attn_weights = torch.bmm(query_states, key_states.transpose(1, 2)) # (bs * 8, seq_len_img, seq_len_text) | |
| 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.stable_softmax_2d: | |
| attn_weights = attn_weights - attn_weights.max() | |
| 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 | |
| attn_weights_T = attn_weights.transpose(1, 2) | |
| attn_weights_l = (attn_weights_T - torch.max(attn_weights_T, dim=-1, keepdim=True)[ | |
| 0]) | |
| if self.clamp_min_for_underflow: | |
| attn_weights_l = torch.clamp(attn_weights_l, min=-50000) # Do not increase -50000, data type half has quite limited range | |
| if self.clamp_max_for_overflow: | |
| attn_weights_l = torch.clamp(attn_weights_l, max=50000) # Do not increase 50000, data type half has quite limited range | |
| attn_weights_l = attn_weights_l.softmax(dim=-1) | |
| # assert attention_mask_l.dtype == torch.int64 | |
| if attention_mask_l is not None: | |
| assert (attention_mask_l.dim() == 2) # (bs, seq_len) | |
| attention_mask = attention_mask_l.unsqueeze(1).unsqueeze(1) # (bs, 1, 1, seq_len) | |
| 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_probs_l = F.dropout(attn_weights_l, p=self.dropout, training=self.training) | |
| attn_output_v = torch.bmm(attn_probs_v, value_l_states) | |
| attn_output_l = torch.bmm(attn_probs_l, value_v_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()}" | |
| ) | |
| if attn_output_l.size() != (bsz * self.num_heads, src_len, self.head_dim): | |
| raise ValueError( | |
| f"`attn_output_l` should be of size {(bsz, self.num_heads, src_len, self.head_dim)}, but is {attn_output_l.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_l = attn_output_l.view(bsz, self.num_heads, src_len, self.head_dim) | |
| attn_output_l = attn_output_l.transpose(1, 2) | |
| attn_output_l = attn_output_l.reshape(bsz, src_len, self.embed_dim) | |
| attn_output_v = self.out_v_proj(attn_output_v) | |
| attn_output_l = self.out_l_proj(attn_output_l) | |
| return attn_output_v, attn_output_l | |
| class BiAttentionBlockForCheckpoint(nn.Module): | |
| def __init__(self, v_dim, l_dim, embed_dim, num_heads, dropout=0.1, | |
| drop_path=.0, init_values=1e-4, ): | |
| """ | |
| Inputs: | |
| embed_dim - Dimensionality of input and attention feature vectors | |
| 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(BiAttentionBlockForCheckpoint, self).__init__() | |
| # pre layer norm | |
| self.layer_norm_v = nn.LayerNorm(v_dim) | |
| self.layer_norm_l = nn.LayerNorm(l_dim) | |
| self.attn = BiMultiHeadAttention(v_dim=v_dim, | |
| l_dim=l_dim, | |
| embed_dim=embed_dim, | |
| num_heads=num_heads, | |
| dropout=dropout, | |
| ) | |
| # 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) | |
| self.gamma_l = nn.Parameter(init_values * torch.ones((l_dim)), requires_grad=True) | |
| def forward(self, v, l, attention_mask_l=None, task=None): | |
| # v: visual features, (bs, sigma(HW), 256) | |
| # l: language features, (bs, seq_len, 768) | |
| v = self.layer_norm_v(v) | |
| l = self.layer_norm_l(l) | |
| delta_v, delta_l = self.attn(v, l, attention_mask_l=attention_mask_l) | |
| # v, l = v + delta_v, l + delta_l | |
| v = v + self.drop_path(self.gamma_v * delta_v) | |
| l = l + self.drop_path(self.gamma_l * delta_l) | |
| return v, l | |