Deformable DETR 论文精读 + 代码详解

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

def get_reference_points(spatial_shapes, valid_ratios, device):
  reference_points_list = []
  for lvl, (H_, W_) in enumerate(spatial_shapes):
      # 从0.5到H-0.5采样H个点，W同理 这个操作的目的也就是为了特征图的对齐
      ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                      torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
      ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
      ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
      ref = torch.stack((ref_x, ref_y), -1)
      reference_points_list.append(ref)
  reference_points = torch.cat(reference_points_list, 1)
  reference_points = reference_points[:, :, None] * valid_ratios[:, None]
  return reference_points

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

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 # batch size, number token, number head, head dims
    # Lq_: number query, L_: level number, P_: sampling number采样点数
    _, Lq_, M_, L_, P_, _ = sampling_locations.shape
    # 按照level划分value
    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
    # [0, 1] -> [-1, 1] 因为要满足F.grid_sample的输入要求
    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_
        # 用双线性插值从feature map上获取value，因为mask的原因越界所以要zeros的方法进行填充
        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_)
    # 不同scale计算出的multi head attention 进行相加，返回output后还需要过一个Linear层
    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()

1
2
3
4
5

# scale-level embedding
# 对4个特征层每层附加256-dim的embedding
# 目的是为了区分query对应到哪个特征层，它会与position embedding相加在一起
# 注意: 位于同一个特征的所有query都会对应到相同的scale-level embedding
self.level_embed = nn.Parameter(torch.Tensor(num_feature_levels, d_model))

1
2
3
4
5
6
7

# 由于 position embedding仅区分h, w的位置，因此对于不同特征层有相同坐标值的特征点来说，是无法区分的，于是这里附加上scale-level embedding作为特征层的区分信息，这样，所有特征点的位置信息就各不相同了
# (bs, c, h, w) => (bs, h*w, c)
pos_embed = pos_embed.flatten(2).transpose(1, 2)
# (bs, h*w, c) + (1, 1, 256)
# note that c = 256 here

lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

class MSDeformAtten(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 dimensions
    :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_heads to a power of 2 which is more efficient in out 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",
      "power of 2",
      "which is more efficient in out CUDA implementation.")
    # 用于cuda实现
    self.im2col_step = 64
    self.d_model = d_model
    self.n_levels = n_levels # 4
    self.n_heads = n_heads
    self.n_points = n_points # 4

    # 采样点的坐标偏移， 每个query在每个注意力头和每个特征层都需要采样n_points个
    # 由于x, y坐标都有对应的偏移量，因此还要*2
    self.sampling_offsets =  nn.Linear(d_model, n_heads * n_levels * n_points * 2)
    # 每个query对应的所有采样点的注意力权重
    self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
    # 线性变换得到value
    self.value_proj = nn.Linear(d_model, d_model)
    # 最后经过这个线性变换得到输出结果
    self.output_proj = nn.Liear(d_model, d_model)

    self.__reset_parameters()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def _reset_parameters(self):
  constant_(self.sampling_offsets.weight.data, 0)
  """初始化偏移量预测的偏置(bias), 使得初始偏移位置犹如不同大小的方形卷积核组合"""

  # (8,) [0, pi / 4, pi / 2, 3 * pi / 2, ..., 7 * pi / 4]
  thetas = torch.arange(self.n_heads, dtype = torch.float32).(2.0 * math.pi / self.n_heads)
  # (8, 2)
  grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
  # grid_init / grid_init.abs().max(-1, keepdi=True)[0]这步计算得到8个头对应的坐标偏移:
  # (1, 0), (1, 1), (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1)
  # 从图形视觉上来看， 形成的偏移位置相当于是3x3， 5x5, 7x7, 9x9正方形卷积核(出去中心，中心是参考点本身)
  for i in range(self.n_points):
    grid_init[:, :, i, :] *= i + 1
  # 注意这里取消了梯度，只是借助nn.Parameter把数值设置进去
  with torch.no_grad():
    self.sampling_offsets.bias = nn.Parameters(grid_init.view(-1))

1
2
3
4
5
6

constant_(self.attention_weights.weight.data, 0)
# TODO: 这里与paper描述的有出入， paper中说bias初始化为1/LK, 其中L为特征层数=4， K为每层的采样点数量=4
# paper中说的那样才是在多有采样点之间归一化
# 否则， 这里所示的， weight和bias都是0， 直接导致最终的输出全为0
cosntant_(self.attention_weights.bias.data, 0)
# constant_(self.attention_weights.bias.data, 1/(self.n_levels * self.n_points))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

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), range in
                                  [0, 1], top-left (0, 0), bottom-right (1, 1)
                                  including padding area or (N, Length_{query}, n_levels, 4),
                                  add additional (w, h) to form reference boxes
  :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, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
  :param input_padding_mask       (N, \sum_{l=0}^{L-1} H_l \cdot W_l), True for padding elements, False for non-padding elements

  :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)


  # 以下主要是计算出采样点的位置。2-stage模式下，输入到Decoder的参考点是Encoder预测的top-k proposal boxes，也就是说是4d的(非2-stage情况下是2d)，于是需要分情况处理:
  sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
  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)
      sampling_locations = reference_points[:, :, None, :, None, :] \
                            + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
  elif reference_points.shape[-1] == 4:
    # reference_points 最后一维中的前两个事中心坐标xy， 后两个是宽高wh
    # 由于初始化时，offset的在-k~k(k = n_points)范围，因此这里除以n_points相当于归一化到0~1
    # 然后乘以宽和高的一半， 加上参考点的中心坐标，这样就是的偏移后采样点位于proposal bbox内
    # 相当于对采样范围进行了约束，减小了搜索空间
      sampling_locations = reference_points[:, :, None, :, None, :2] \
                            + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
  else:
      raise ValueError(
          'Last dim of reference_points must be 2 or 4, but get {} instead.'.format(reference_points.shape[-1]))
  # 根据采样点位置拿出对应的value，并且施加预测出来的注意力权重(和value进行weighted sum)
  # (N, Len_in, 256)
  # 注意: 实质调用的是基于CUDA实现的版本，需要编译
  output = MSDeformAttnFunction.apply(
      value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)
  # (N， Len_in, 256)
  output = self.output_proj(output)
  return output

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

def ms_deform_attn_core_pytorch(value, value_spatial_shapes, sampling_locations, attention_weights):
  """多尺度可变性注意力， 根据采样点的位置在多尺度value中插值采样出对应的特征图，最后和注意力权重进行weighted sum得到输出"""
  # for debug and test only,
  # need to use cuda version instead
  N_, S_, M_, D_ = value.shape
  _, Lq_, M_, L_, P_, _ = sampling_locations.shape
  # 由于以下使用了F.grid_sample()，要求采样位置的坐标是归一化到[-1, 1] ((-1, -1)代表左上角， (1，1)代表右下角)
  # 因此， 这里是将[0, 1]映射到[-1, 1]
  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)
      # 根据采样点坐标在value中插值出对应的特征
      # ps: grid_sample()用法
      # 这里value_l 充当被插值采样的特征图，是input， 维度需要时 4D/5D
      # sampling_grid_l则代表采样的位置，是grid，最后一维2对应input中的坐标(可能是小数)
      # 倒数第2，3维代表采样后输出特征图宽、高
      # input和grid的第一个维度必须一致，最终输出的通道数与input一致，是不变的
      # 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_)
  # 最后就是将注意力权重和采样特征进行weighted sum:
  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()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

    # prepare input for decoder
    # c = 256 中间那一维等于(所有层)特征点的数量
    bs, _, c = memory.shape
    # 根据是否2-stage分情况进行处理，因为生成的reference points不同
    if self.two_stage:
      # 生成proposals， 并且对Encoder的输出(memory)进行处理(全连接层 + 归一化)
      #(bs, h_lvl1 * w_lvl1 + h_lvl2 * w_lvl2 + ..., 256), (bs, h_lvl1 * w_lvl1 + h_lvl2 * w_lvl2 + ..., 4)
      # 其中proposals每个都是xywh形式， 并且是经过inverse-sigmoid函数后的结果
      # (其实这里的output_proposals对应的就是各层特征图各个特征点的位置(相当于anchor的形式，是固定的)，
      # 因此还需要借助Decoder最后一层的bbox head来预测一个偏移(offset)来得到一个更加灵活的结果，
      # 这才是第一阶段预测的proposal boxes)
      output_memory, output_proposals = self.gen_encoder_output_proposals(memory, mask_flatten, spatial_shapes)

      # hack implementation for two-stage Deformable DETR
      # 注意: 这里维度对应的是多分类，并非二分类
      enc_outputs_class = self.decoder.class_embed[self.decoder.num_layers](output_memory)
      # bbox head预测的是相对proposals的偏移，因此这里要相加， 后续还要经过sigmoid函数才得到真正的bbox预测结果(归一化形式)
      #(bs， h_lvl1 * w_lvl1 + h_lvl2 * w_lvl2 + .., 4)
      enc_outputs_coord_unact = self.decoder.bbox_embed[self.decoder.num_layers](output_memory) + output_proposals

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

      # 300
      topk = self.two_stage_num_proposals
      # 选取得分最高的top分类预测，最后的[1]代表取得返回top对应的索引
      # (bs， k = 300)
      # TODO: 取第一个类别的预测结果算top-k，代表二分类
      # 当不适用iterative bbox refine时， 所有class_embed参数共享，这样会使得在第二阶段对解码输出进行分类时都偏向于第一个类别
      # 这样貌似不妥
      topk_proposals = torch.topk(enc_outputs_class[..., 0], topk, dim=1)[1]
      # 拿出top-k得分最高的对应的预测bbox:  (bs, k = 300, 4)
      topk_coords_unact = torch.gather(enc_outputs_coord_unact, 1, topk_proposals.unsqueeze(-1).repeat(1, 1, 4))
      # 注意: 这里取消了梯度
      topk_coords_unact = topk_coords_unact.detach()
      # 经过sigmoid，变成了归一化形式，这个结果会送到decoder中作为初始的bboxes估计
      reference_points = topk_coords_unact.sigmoid()
      # init_reference_out = reference_points
      init_reference_out = reference_points

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# 如果使用2-stage， 那么在Decoder 中多加一层(变为7层)，用于第一阶段中proposal的预测输出
# (预测输出实际上由Encoder输出，只不过这里借用一层Decoder来解码形成预测结果，也就是借用了分类和回归预测的头部)
num_pred = (transformer.decoder.num_layers + 1) if two_stage else transformer.decoder.num_layers
"""在iterative box refine策略下，_get_clones得到的每个模块都是不同的实例，参数不共享；
  而不使用该策略时， nn.ModuleList中每个都是相同的实例，即参数共享"""
# TODO: 以下bbox head的bias的后两个初始化为-2.0是为何？
if with_box_refine:
  # 如果使用迭代的bbox校正策略，则Decoder各层参数不共享，因此这里用_get_clones(deepcopy)

  self.class_embed = _get_clones(self.class_embed, num_pred)
  self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
  nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
  # hack implementation for iterative bounding box refinement
  # 默认情况下， Decoder的bbox_embed设置为None，因此只有使用iterative bbox refinement策略时，
  # 其bbox_embed才不是None； 在使用iterative bbox refine时，Decoder每层都会预测bbox偏移量，
  # 使用这个偏移量对上一层的预测输出进行校正
  self.transformer.decoder.bbox_embed = self.bbox_embed
else:
  nn.init.contant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
  # 不使用iterative bbox refine 策略，则各层参数共享
  self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
  self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
  # 不用iterative bbox refine策略， 则Decoder的bbox_embed设置为None
  self.transofmer.decoder.bbox_embed = None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

  reference_points = topk_coords_unact.sigmoid()
  # init_reference_out = reference_points
  init_reference_out = reference_points
  """
  生成的Decoder的query(target)和query embedding:
    - 对于top-k proposal boxes进行位置编码，编码方式是给xywh每个都赋予128维，
    其中每128维中，偶数维度用sin函数，奇数维度用cos函数编码；
    然后经过全连接层和归一化处理；
    最终， 前256维结果对应xy作为Decoder 的query embedding(因为xy代表的是位置信息)
  """
  #(bs, k = 300, 4 x 128 = 512)
  pos_trans_out = self.pos_trans_norm(self.pos_trans(self.get_proposal_pos_embed(topk_coords_unact)))
  # (bs, k = 300, 256), (bs, k = 300, 256)
  query_embed, tgt = torch.split(pos_trans_out, c, dim=2)
else:
  # 仅为了与2-stage的情况兼容
  enc_output_class = enc_outputs_coord_unact = None
  # (n_query = 300, 256) (n_query = 300, 256)
  query_embed, tgt = torch.split(query_embed, c, dim=1)
  # (bs, n_query = 300, 256)
  query_embed = query_embed.unsqueeze(0).expand(bs, -1, -1)
  # (bs, n_query = 300, 256)
  tgt = tgt.unsqueeze(0).expand(bs, -1, -1)
  # 通过全连接层生成proposal参考点的归一化坐标(cx, cy) : (bs, n_query = 300, 2)
  reference_points = self.reference_points(query_embed).sigmoid()
  init_reference_out = reference_points

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

  # decoder
  # hs: (n_dec_layers, bs, n_query=300, d_model=256);
  # itner_references: with iterative bbox refine - (n_dec_layers, bs, k = 300, 4)
  # otherwise - (n_dec_layers, bs, k = 300, 2)
  hs, inter_references = self.decoder(tgt, reference_points, memory,
                                      spatial_shapes, level_start_index, valid_ratios, query_embed, mask_flatten)

  inter_references_out = inter_references
  if self.two_stage:
      return hs, init_reference_out, inter_references_out, enc_outputs_class, enc_outputs_coord_unact
  # 1-stage情况下:  enc_outputs_class 和 enc_outputs_coord_unact 都是None
  return hs, init_reference_out, inter_references_out, None, None

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
  # 这里的pos是position embedding + scale-level embedding
  output = src
  # 将各特征点在其所在特征层的归一化坐标映射到所有特征层，使得每个特征点在所有特征层上都会得到一个归一化的坐标
  # 这个reference_points 相当于key的角色， 从而每个query都会和其在所有特征层的位置(也就是以下计算出来的坐标)进行交互
  # 实现了跨尺度融合的效果，因此不需要FPN
  # (bs, H_lvl1 * W_lvl1 + H_lvl2 * W_lvl2 + ..., n_lvl, 2)
  reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
  for _, layer in enumerate(self.layers):
      output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)

  return output

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

def get_reference_points(spatial_shapes, valid_ratios, device):
  reference_points_list = []
  for lvl, (H_, W_) in enumerate(spatial_shapes):
    # (H_, W_), (H_, W_)
    # 0.5 是为对应到特征点中心
    ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                  torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
    """将各层特征图每个特征点中心坐标根据特征图非padding的边长进行归一化(可能大于1)"""
    # (1, H_*W_) / (bs, 1) 后一项是特征图有效(非padding)部分的高
    ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
    # (1, H_*W_) / (bs, 1) 后一项是特征图有效(非padding)部分的宽
    ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
    # (bs, H_*W_, 2) 每一项是xy
    ref = torch.stack((ref_x, ref_y), -1)
    reference_points_list.append(ref)
  # (bs, H_lvl1 * W_lvl1 + H_lvl2 * W_lvl2 + ..., 2)
  reference_points = torch.cat(reference_points_list, 1)
  """将各特征点在其所在特征层的归一化坐标映射(扩散)到所有特征层， 这样每个特征点在所有特征层上都会得到一个归一化坐标"""
  # TODO: 以下这样貌似不妥， 如果各特征层对应的valid_ratio不一致，
  # TODO: 则坐标值有可能大于1， 而后续没有再对这里的reference_points进行归一化到 0~1
  # (bs, H_lvl1 * W_lvl1 + H_lvl2 * W_lvl2 + ..., 1, 2) * (bs, 1, n_lvl, 2)
  reference_points = reference_points[:, :, None] * valid_ratios[:, None]
  # (bs, H_lvl1 * W_lvl1 + H_lvl2 * W_lvl2 + ..., n_lvl, 2)
  return reference_points

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
  """Deformable DETR的Encoder也是由self-attention + FNN组成
    只不过这里self-attention使用Multi-Scale Deformable Attention， 并且位置编码加入了scale-level embedding
  """
  # self attention
  # 这里的pos是position embedding + scale-level embedding
  # padding_mask 就是指示各特征图哪些位置是原图padding的
  # reference_points 就是每个特征点本身中心的位置(归一化坐标): (bs, H_lvl1 * W_lvl1 + H_lvl2 * W_lvl2 + ..., n_lvl, 2)
  # 注意一个特征点不仅在其所有特征层有个坐标， 而且还在其他特征层也都分别映射了一个坐标
  # self attention
  src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
  src = src + self.dropout1(src2)
  src = self.norm1(src)

  # ffn
  src = self.forward_ffn(src)

  return src

1
2
3
4

# 每一层输入的output是上一层输出的结果，而reference_points_input在使用iterative bbox refine策略时，
# 每层都会对齐进行校正， 因此下一层用到的也是上一层的输出结果
# (bs， n_query=300，hidden_dim=256)
output = layer(output, query_pos, reference_points_input, src, src_spatial_shapes, src_level_start_index, src_padding_mask)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

def forward(self, tgt, reference_points, src, src_spatial_shapes, src_level_start_index, src_valid_ratios,
                query_pos=None, src_padding_mask=None):
  # 说明一下Decoder一开始得到的tgt, query_pos和reference_points, 分为两种情况:
  # 1. 2-stage 模式下，reference_points 是Encoder输出的top-k proposal boxes(并归一化)，最后一维为4
  # 而tgt和query_pos由其经过position embedding得到；
  # 2. 1-stage 模式下， tgt和query_pos是预设的embedding， reference_points通过这个query_pos经全连接层得到，
  # 最后一维为2
  # 另外，src是Encoder最终编码输出的特征图，即 memory
  output = tgt

  # 中间各层(包括头尾)的解码输出
  intermediate = []
  # 中间各层(包括头尾)校正的参考点
  intermediate_reference_points = []
  for lid, layer in enumerate(self.layers):
    # 2-stage 模式下， 参考点是proposal boxes， 因此最后一维是4
    if reference_points.shape[-1] == 4:
      # (bs, k = 300, n_lvl, 4)
      # (bs, k = 300, 1, 4) * (bs, 1, n_lvl, 4)
      reference_points_input = reference_points[:, :, None] \
                                * torch.cat([src_valid_ratios, src_valid_ratios], -1)[:, None]
    else:
      # 1-stage 模式下 参考点就是通过query embedding 变换而来的中心坐标形式，因此最后一维是2
      assert reference_points.shape[-1] == 2
      # (bs, k=300, n_lvl, 2)
      # (bs, k=300, 1, 2) * (bs, 1, n_lvl, 2)
      reference_points_input = reference_points[:, :, None] * src_valid_ratios[:, None]
    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](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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def forward(self, tgt, query_pos, reference_points, src, src_spatial_shapes, level_start_index, src_padding_mask=None):
  # 若是2-stage， 则tgt 和 query_pos来自Encoder输出的top-k proposal boxes(经过位置嵌入)
  # 而reference_points 就是这个top-k proposal boxes(归一化)
  # 否则， tgt和query_pos由预设的embedding产生， 而reference_points由query_pos经过全连接层生成

  # self attention
  # (bs, k = 300, d_model=256)
  q = k = self.with_pos_embed(tgt, query_pos)
  # (bs, k = 300, d_model=256) 注意: value就是target本身不需要， 不需要位置编码
  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)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

# 上续 decoder的 decoder layer forward
# cross attention
# src是Encoder输出的memory， 即编码后的特征(bs, n_feat_points, d_model=256), 其会经过线性变换得到value,
# 这里的tgt来自self-attention的输出，而query_pos依旧如刚传进来Decoder时一样，不变
# reference_points: (bs, k=300, n_feat_lvl, 4 or 2) 在cross-attention中代表key的位置信息
tgt2 = self.cross_attn(self.with_pos_embed(tgt, query_pos),
                        reference_points,
                        src, src_spatial_shapes, level_start_index, src_padding_mask)
tgt = tgt + self.dropout1(tgt2)
tgt = self.norm1(tgt)

1
2
3
4

# ffn
tgt = self.forward_ffn(tgt)

return tgt

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# hack implementation for iterative bounding box refinement
# 当使用了iterative bbox refine策略，则这里的bbox_embed就不是None
# 并且会对reference points进行refine， 之后每层的reference points都是前一层校正后的结果(取消了梯度)
# 否则，即没有使用iterative bbox refine的话， 那么reference points将永远是一样的
# (来自Encoder输出的proposal boxes或由预设的embedding通过位置编码，需要根据是2-stage还是1-stage的情况而定)
if self.bbox_embed is not None:
  # (bs， n_query=300, 4)
  tmp = self.bbox_embed[lid](output)
  # 2-stage 模式
  if reference_points.shape[-1] == 4:
    # tmp是bbox head输出的(相对参考点也就是proposal boxes)预测偏移量
    new_reference_points = tmp + inverse_sigmoid(reference_points)
    # (bs，k=300， 4)
    new_reference_points = new_reference_points.sigmoid()
  else:
    assert reference_points.shape[-1] == 2
    # 注意: 即使是1-stage， 在iterative bbox refine策略下，这里
    # 也将reference points最后一维变成4
    new_reference_points = tmp
    # 1-stage模式下，参考点是特征点中心坐标(最后一维是2)， 因此这里预测的偏移量只需要去前面两维做加法
    new_reference_points[..., :2] = tmp[..., :2] + inverse_sigmoid(reference_points)
    # (bs， k=300, 4)
    new_reference_points = new_reference_points.sigmoid()
  # (bs, k=300, 4)
  # 注意:  这里取消了梯度！
  reference_points = new_reference_points.detach()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

# DECODER forward()
  if self.return_intermediate:
    intermediate.append(output)
    intermediate_reference_points.append(reference_points)

if self.return_intermediate:
  # (n_layers, bs, n_query=300, d_model), (n_layers, bs, k=300, 4 or 2)
  return torch.stack(intermediate), torch.stack(intermediate_reference_points)
# 这里应该是[output] [reference_points] 这样才兼容return_intermediate的情况，第一个维度对应Decoder的层数
# return output, reference_points
return output, reference_points