# The implementation here is modified based on timm, # originally Apache 2.0 License and publicly available at # https://github.com/naver-ai/vidt/blob/vidt-plus/methods/vidt/deformable_transformer.py import copy import math import warnings import torch import torch.nn.functional as F from timm.models.layers import DropPath from torch import nn from torch.nn.init import constant_, normal_, xavier_uniform_ class DeformableTransformer(nn.Module): """ A Deformable Transformer for the neck in a detector The transformer encoder is completely removed for ViDT Args: d_model: the channel dimension for attention [default=256] nhead: the number of heads [default=8] num_decoder_layers: the number of decoding layers [default=6] dim_feedforward: the channel dim of point-wise FFNs [default=1024] dropout: the degree of dropout used in FFNs [default=0.1] activation: An activation function to use [default='relu'] return_intermediate_dec: whether to return all the indermediate outputs [default=True] num_feature_levels: the number of scales for extracted features [default=4] dec_n_points: the number of reference points for deformable attention [default=4] drop_path: the ratio of stochastic depth for decoding layers [default=0.0] token_label: whether to use the token label loss for training [default=False]. This is an additional trick proposed in https://openreview.net/forum?id=LhbD74dsZFL (ICLR'22) for further improvement """ def __init__(self, d_model=256, nhead=8, num_decoder_layers=6, dim_feedforward=1024, dropout=0.1, activation='relu', return_intermediate_dec=True, num_feature_levels=4, dec_n_points=4, drop_path=0., token_label=False): super().__init__() self.d_model = d_model self.nhead = nhead decoder_layer = DeformableTransformerDecoderLayer( d_model, dim_feedforward, dropout, activation, num_feature_levels, nhead, dec_n_points, drop_path=drop_path) self.decoder = DeformableTransformerDecoder(decoder_layer, num_decoder_layers, return_intermediate_dec) self.level_embed = nn.Parameter( torch.Tensor(num_feature_levels, d_model)) self.token_label = token_label self.reference_points = nn.Linear(d_model, 2) if self.token_label: self.enc_output = nn.Linear(d_model, d_model) self.enc_output_norm = nn.LayerNorm(d_model) self.token_embed = nn.Linear(d_model, 91) prior_prob = 0.01 bias_value = -math.log((1 - prior_prob) / prior_prob) self.token_embed.bias.data = torch.ones(91) * bias_value self._reset_parameters() def _reset_parameters(self): for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) for m in self.modules(): if isinstance(m, MSDeformAttn): m._reset_parameters() normal_(self.level_embed) def get_proposal_pos_embed(self, proposals): num_pos_feats = 128 temperature = 10000 scale = 2 * math.pi dim_t = torch.arange( num_pos_feats, dtype=torch.float32, device=proposals.device) dim_t = temperature**(2 * (dim_t // 2) / num_pos_feats) # N, L, 4 proposals = proposals.sigmoid() * scale # N, L, 4, 128 pos = proposals[:, :, :, None] / dim_t # N, L, 4, 64, 2 pos = torch.stack((pos[:, :, :, 0::2].sin(), pos[:, :, :, 1::2].cos()), dim=4).flatten(2) return pos def gen_encoder_output_proposals(self, memory, memory_padding_mask, spatial_shapes): N_, S_, C_ = memory.shape proposals = [] _cur = 0 for lvl, (H_, W_) in enumerate(spatial_shapes): mask_flatten_ = memory_padding_mask[:, _cur:(_cur + H_ * W_)].view( N_, H_, W_, 1) valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1) valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1) grid_y, grid_x = torch.meshgrid( torch.linspace( 0, H_ - 1, H_, dtype=torch.float32, device=memory.device), torch.linspace( 0, W_ - 1, W_, dtype=torch.float32, device=memory.device)) grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1) scale = torch.cat([valid_W.unsqueeze(-1), valid_H.unsqueeze(-1)], 1).view(N_, 1, 1, 2) grid = (grid.unsqueeze(0).expand(N_, -1, -1, -1) + 0.5) / scale wh = torch.ones_like(grid) * 0.05 * (2.0**lvl) proposal = torch.cat((grid, wh), -1).view(N_, -1, 4) proposals.append(proposal) _cur += (H_ * W_) output_proposals = torch.cat(proposals, 1) tmp = (output_proposals > 0.01) & (output_proposals < 0.99) output_proposals_valid = tmp.all(-1, keepdim=True) output_proposals = torch.log(output_proposals / (1 - output_proposals)) output_proposals = output_proposals.masked_fill( memory_padding_mask.unsqueeze(-1), float('inf')) output_proposals = output_proposals.masked_fill( ~output_proposals_valid, float('inf')) output_memory = memory output_memory = output_memory.masked_fill( memory_padding_mask.unsqueeze(-1), float(0)) output_memory = output_memory.masked_fill(~output_proposals_valid, float(0)) output_memory = self.enc_output_norm(self.enc_output(output_memory)) return output_memory, output_proposals def get_valid_ratio(self, mask): _, H, W = mask.shape valid_H = torch.sum(~mask[:, :, 0], 1) valid_W = torch.sum(~mask[:, 0, :], 1) valid_ratio_h = valid_H.float() / H valid_ratio_w = valid_W.float() / W valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1) return valid_ratio def forward(self, srcs, masks, tgt, query_pos): """ The forward step of the decoder Args: srcs: [Patch] tokens masks: input padding mask tgt: [DET] tokens query_pos: [DET] token pos encodings Returns: hs: calibrated [DET] tokens init_reference_out: init reference points inter_references_out: intermediate reference points for box refinement enc_token_class_unflat: info. for token labeling """ # prepare input for the Transformer decoder src_flatten = [] mask_flatten = [] spatial_shapes = [] for lvl, (src, mask) in enumerate(zip(srcs, masks)): bs, c, h, w = src.shape spatial_shape = (h, w) spatial_shapes.append(spatial_shape) src = src.flatten(2).transpose(1, 2) mask = mask.flatten(1) src_flatten.append(src) mask_flatten.append(mask) src_flatten = torch.cat(src_flatten, 1) mask_flatten = torch.cat(mask_flatten, 1) spatial_shapes = torch.as_tensor( spatial_shapes, dtype=torch.long, device=src_flatten.device) level_start_index = torch.cat((spatial_shapes.new_zeros( (1, )), spatial_shapes.prod(1).cumsum(0)[:-1])) valid_ratios = torch.stack([self.get_valid_ratio(m) for m in masks], 1) memory = src_flatten bs, _, c = memory.shape tgt = tgt # [DET] tokens query_pos = query_pos.expand(bs, -1, -1) # [DET] token pos encodings # prepare input for token label if self.token_label: output_memory, output_proposals = self.gen_encoder_output_proposals( memory, mask_flatten, spatial_shapes) enc_token_class_unflat = None if self.token_label: enc_token_class = self.token_embed(output_memory) enc_token_class_unflat = [] for st, (h, w) in zip(level_start_index, spatial_shapes): enc_token_class_unflat.append( enc_token_class[:, st:st + h * w, :].view(bs, h, w, 91)) # reference points for deformable attention reference_points = self.reference_points(query_pos).sigmoid() init_reference_out = reference_points # query_pos -> reference point # decoder hs, inter_references = self.decoder(tgt, reference_points, memory, spatial_shapes, level_start_index, valid_ratios, query_pos, mask_flatten) inter_references_out = inter_references return hs, init_reference_out, inter_references_out, enc_token_class_unflat class DeformableTransformerDecoderLayer(nn.Module): """ A decoder layer. Args: d_model: the channel dimension for attention [default=256] d_ffn: the channel dim of point-wise FFNs [default=1024] dropout: the degree of dropout used in FFNs [default=0.1] activation: An activation function to use [default='relu'] n_levels: the number of scales for extracted features [default=4] n_heads: the number of heads [default=8] n_points: the number of reference points for deformable attention [default=4] drop_path: the ratio of stochastic depth for decoding layers [default=0.0] """ def __init__(self, d_model=256, d_ffn=1024, dropout=0.1, activation='relu', n_levels=4, n_heads=8, n_points=4, drop_path=0.): super().__init__() # [DET x PATCH] deformable cross-attention self.cross_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points) self.dropout1 = nn.Dropout(dropout) self.norm1 = nn.LayerNorm(d_model) # [DET x DET] self-attention self.self_attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout) self.dropout2 = nn.Dropout(dropout) self.norm2 = nn.LayerNorm(d_model) # ffn for multi-heaed self.linear1 = nn.Linear(d_model, d_ffn) self.activation = _get_activation_fn(activation) self.dropout3 = nn.Dropout(dropout) self.linear2 = nn.Linear(d_ffn, d_model) self.dropout4 = nn.Dropout(dropout) self.norm3 = nn.LayerNorm(d_model) # stochastic depth self.drop_path = DropPath(drop_path) if drop_path > 0. else None @staticmethod def with_pos_embed(tensor, pos): return tensor if pos is None else tensor + pos def forward_ffn(self, tgt): tgt2 = self.linear2(self.dropout3(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout4(tgt2) tgt = self.norm3(tgt) return tgt def forward(self, tgt, query_pos, reference_points, src, src_spatial_shapes, level_start_index, src_padding_mask=None): # [DET] self-attention q = k = self.with_pos_embed(tgt, query_pos) tgt2 = self.self_attn( q.transpose(0, 1), k.transpose(0, 1), tgt.transpose(0, 1))[0].transpose(0, 1) tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt) # Multi-scale deformable cross-attention in Eq. (1) in the ViDT paper tgt2 = self.cross_attn( self.with_pos_embed(tgt, query_pos), reference_points, src, src_spatial_shapes, level_start_index, src_padding_mask) if self.drop_path is None: tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt) # ffn tgt = self.forward_ffn(tgt) else: tgt = tgt + self.drop_path(self.dropout1(tgt2)) tgt2 = self.linear2( self.dropout3(self.activation(self.linear1(tgt)))) tgt = tgt + self.drop_path(self.dropout4(tgt2)) tgt = self.norm3(tgt) return tgt class DeformableTransformerDecoder(nn.Module): """ A Decoder consisting of multiple layers Args: decoder_layer: a deformable decoding layer num_layers: the number of layers return_intermediate: whether to return intermediate results """ def __init__(self, decoder_layer, num_layers, return_intermediate=False): super().__init__() self.layers = _get_clones(decoder_layer, num_layers) self.num_layers = num_layers self.return_intermediate = return_intermediate # hack implementation for iterative bounding box refinement self.bbox_embed = None self.class_embed = None def forward(self, tgt, reference_points, src, src_spatial_shapes, src_level_start_index, src_valid_ratios, query_pos=None, src_padding_mask=None): """ The forward step of the Deformable Decoder Args: tgt: [DET] tokens reference_points: reference points for deformable attention src: the [PATCH] tokens fattened into a 1-d sequence src_spatial_shapes: the spatial shape of each multi-scale feature map src_level_start_index: the start index to refer different scale inputs src_valid_ratios: the ratio of multi-scale feature maps query_pos: the pos encoding for [DET] tokens src_padding_mask: the input padding mask Returns: output: [DET] tokens calibrated (i.e., object embeddings) reference_points: A reference points If return_intermediate = True, output & reference_points are returned from all decoding layers """ output = tgt intermediate = [] intermediate_reference_points = [] # iterative bounding box refinement (handling the [DET] tokens produced from Swin with RAM) if self.bbox_embed is not None: tmp = self.bbox_embed[0](output) if reference_points.shape[-1] == 4: new_reference_points = tmp + inverse_sigmoid(reference_points) new_reference_points = new_reference_points.sigmoid() else: assert reference_points.shape[-1] == 2 new_reference_points = tmp new_reference_points[ ..., :2] = tmp[..., :2] + inverse_sigmoid(reference_points) new_reference_points = new_reference_points.sigmoid() reference_points = new_reference_points.detach() # if self.return_intermediate: intermediate.append(output) intermediate_reference_points.append(reference_points) for lid, layer in enumerate(self.layers): if reference_points.shape[-1] == 4: tmp0 = reference_points[:, :, None] tmp1 = torch.cat([src_valid_ratios, src_valid_ratios], -1)[:, None] reference_points_input = tmp0 * tmp1 else: assert reference_points.shape[-1] == 2 reference_points_input = reference_points[:, :, None] * src_valid_ratios[:, None] # deformable operation output = layer(output, query_pos, reference_points_input, src, src_spatial_shapes, src_level_start_index, src_padding_mask) # hack implementation for iterative bounding box refinement if self.bbox_embed is not None: tmp = self.bbox_embed[lid + 1](output) if reference_points.shape[-1] == 4: new_reference_points = tmp + inverse_sigmoid( reference_points) new_reference_points = new_reference_points.sigmoid() else: assert reference_points.shape[-1] == 2 new_reference_points = tmp new_reference_points[..., :2] = tmp[ ..., :2] + inverse_sigmoid(reference_points) new_reference_points = new_reference_points.sigmoid() reference_points = new_reference_points.detach() # if self.return_intermediate: intermediate.append(output) intermediate_reference_points.append(reference_points) if self.return_intermediate: return torch.stack(intermediate), torch.stack( intermediate_reference_points) return output, reference_points def _get_clones(module, N): return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) def _get_activation_fn(activation): """Return an activation function given a string""" if activation == 'relu': return F.relu if activation == 'gelu': return F.gelu if activation == 'glu': return F.glu raise RuntimeError(F'activation should be relu/gelu, not {activation}.') def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights): # for debug and test only, # need to use cuda version instead N_, S_, M_, D_ = value.shape _, Lq_, M_, L_, P_, _ = sampling_locations.shape value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1) sampling_grids = 2 * sampling_locations - 1 sampling_value_list = [] for lid_, (H_, W_) in enumerate(value_spatial_shapes): # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_ value_l_ = value_list[lid_].flatten(2).transpose(1, 2).reshape( N_ * M_, D_, H_, W_) # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2 sampling_grid_l_ = sampling_grids[:, :, :, lid_].transpose(1, 2).flatten(0, 1) # N_*M_, D_, Lq_, P_ sampling_value_l_ = F.grid_sample( value_l_, sampling_grid_l_, mode='bilinear', padding_mode='zeros', align_corners=False) sampling_value_list.append(sampling_value_l_) # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_, M_, 1, Lq_, L_*P_) attention_weights = attention_weights.transpose(1, 2).reshape( N_ * M_, 1, Lq_, L_ * P_) output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights).sum(-1).view(N_, M_ * D_, Lq_) return output.transpose(1, 2).contiguous() def _is_power_of_2(n): if (not isinstance(n, int)) or (n < 0): raise ValueError( 'invalid input for _is_power_of_2: {} (type: {})'.format( n, type(n))) return (n & (n - 1) == 0) and n != 0 class MSDeformAttn(nn.Module): def __init__(self, d_model=256, n_levels=4, n_heads=8, n_points=4): """ Multi-Scale Deformable Attention Module :param d_model hidden dimension :param n_levels number of feature levels :param n_heads number of attention heads :param n_points number of sampling points per attention head per feature level """ super().__init__() if d_model % n_heads != 0: raise ValueError( 'd_model must be divisible by n_heads, but got {} and {}'. format(d_model, n_heads)) _d_per_head = d_model // n_heads # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation if not _is_power_of_2(_d_per_head): warnings.warn( "You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 " 'which is more efficient in our CUDA implementation.') self.im2col_step = 64 self.d_model = d_model self.n_levels = n_levels self.n_heads = n_heads self.n_points = n_points self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2) self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points) self.value_proj = nn.Linear(d_model, d_model) self.output_proj = nn.Linear(d_model, d_model) self._reset_parameters() def _reset_parameters(self): constant_(self.sampling_offsets.weight.data, 0.) thetas = torch.arange( self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads) grid_init = torch.stack([thetas.cos(), thetas.sin()], -1) grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view( self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1) for i in range(self.n_points): grid_init[:, :, i, :] *= i + 1 with torch.no_grad(): self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1)) constant_(self.attention_weights.weight.data, 0.) constant_(self.attention_weights.bias.data, 0.) xavier_uniform_(self.value_proj.weight.data) constant_(self.value_proj.bias.data, 0.) xavier_uniform_(self.output_proj.weight.data) constant_(self.output_proj.bias.data, 0.) def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, input_padding_mask=None): """ :param query (N, Length_{query}, C) :param reference_points (N, Length_{query}, n_levels, 2) :param input_flatten (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l, C) :param input_spatial_shapes (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})] :param input_level_start_index (n_levels, ) :param input_padding_mask (N, \\sum_{l=0}^{L-1} H_l \\cdot W_l) :return output (N, Length_{query}, C) """ N, Len_q, _ = query.shape N, Len_in, _ = input_flatten.shape assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in value = self.value_proj(input_flatten) if input_padding_mask is not None: value = value.masked_fill(input_padding_mask[..., None], float(0)) value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads) sampling_offsets = self.sampling_offsets(query).view( N, Len_q, self.n_heads, self.n_levels, self.n_points, 2) # attn weights for each sampled query. attention_weights = self.attention_weights(query).view( N, Len_q, self.n_heads, self.n_levels * self.n_points) attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points) # N, Len_q, n_heads, n_levels, n_points, 2 if reference_points.shape[-1] == 2: offset_normalizer = torch.stack( [input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1) tmp0 = reference_points[:, :, None, :, None, :] tmp1 = sampling_offsets / offset_normalizer[None, None, None, :, None, :] sampling_locations = tmp0 + tmp1 elif reference_points.shape[-1] == 4: tmp0 = reference_points[:, :, None, :, None, :2] tmp1 = sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5 sampling_locations = tmp0 + tmp1 else: raise ValueError( 'Last dim of reference_points must be 2 or 4, but get {} instead.' .format(reference_points.shape[-1])) output = ms_deform_attn_core_pytorch(value, input_spatial_shapes, sampling_locations, attention_weights) output = self.output_proj(output) return output def inverse_sigmoid(x, eps=1e-5): x = x.clamp(min=0, max=1) x1 = x.clamp(min=eps) x2 = (1 - x).clamp(min=eps) return torch.log(x1 / x2)