在上一篇博客中,博主已经讲解了如何利用Uncertainty-minimal Query Selection
选择出好的特征,接下来便要将这些特征输入到Decoder中进行解码,需要注意的是,在RT-DETR
的Encoder中,使用的是标准的自注意力计算方法,而在其Decoder中,则使用的是可变形自注意力(deformable attention),可变形自注意力能够大幅的降低计算量,同时该部分还使用到了CUDA算子,能够加快运行速度,当然,这个可变形自注意力计算并非是RT-DETR的创新点,但其作用却是极大,在DINO,DN-Deformable-DETR中都有使用。
关于Deformable-DETR,博主曾经介绍过,大家如果有兴趣可以参考博主这篇博文:
Deformable DETR模型学习记录
Decoder参数
输入Decoder的参数如下:
out_bboxes, out_logits = self.decoder(
target,
init_ref_points_unact,
memory,
spatial_shapes,
level_start_index,
self.dec_bbox_head,
self.dec_score_head,
self.query_pos_head,
attn_mask=attn_mask)
target
是查询向量添加噪声以及查询向量筛选后的特征向量,即498=198+30
init_ref_point_unct
是参考点的xywh
(Anchor
)
memory
是Encoder
输出的特征向量
spatial_shapes
是Encoder
输出的三个特征图的维度
记录每个特征图开始的索引(已将特征图展平)
attn_mask
特征图掩膜
query_pos_head
的结构如下:
MLP(
(layers): ModuleList(
(0): Linear(in_features=4, out_features=512, bias=True)
(1): Linear(in_features=512, out_features=256, bias=True)
)
(act): ReLU(inplace=True)
)
TransformerDecoderLayer的结构如下:
ModuleList(
(0-2): 3 x TransformerDecoderLayer(
(self_attn): MultiheadAttention(
(out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
)
(dropout1): Dropout(p=0.0, inplace=False)
(norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(cross_attn): MSDeformableAttention(
(sampling_offsets): Linear(in_features=256, out_features=192, bias=True)
(attention_weights): Linear(in_features=256, out_features=96, bias=True)
(value_proj): Linear(in_features=256, out_features=256, bias=True)
(output_proj): Linear(in_features=256, out_features=256, bias=True)
)
(dropout2): Dropout(p=0.0, inplace=False)
(norm2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(linear1): Linear(in_features=256, out_features=1024, bias=True)
(dropout3): Dropout(p=0.0, inplace=False)
(linear2): Linear(in_features=1024, out_features=256, bias=True)
(dropout4): Dropout(p=0.0, inplace=False)
(norm3): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
)
多层Decoder Layer(TransformerDecoder)
多层DecoderLayer的操作如下:
def forward(self,
tgt,
ref_points_unact,
memory,
memory_spatial_shapes,
memory_level_start_index,
bbox_head,
score_head,
query_pos_head,
attn_mask=None,
memory_mask=None):
output = tgt
dec_out_bboxes = []
dec_out_logits = []
ref_points_detach = F.sigmoid(ref_points_unact)
for i, layer in enumerate(self.layers):
ref_points_input = ref_points_detach.unsqueeze(2)
query_pos_embed = query_pos_head(ref_points_detach)
output = layer(output, ref_points_input, memory,
memory_spatial_shapes, memory_level_start_index,
attn_mask, memory_mask, query_pos_embed)
inter_ref_bbox = F.sigmoid(bbox_head[i](output) + inverse_sigmoid(ref_points_detach))
if self.training:
dec_out_logits.append(score_head[i](output))
if i == 0:
dec_out_bboxes.append(inter_ref_bbox)
else:
dec_out_bboxes.append(F.sigmoid(bbox_head[i](output) + inverse_sigmoid(ref_points)))
elif i == self.eval_idx:
dec_out_logits.append(score_head[i](output))
dec_out_bboxes.append(inter_ref_bbox)
break
ref_points = inter_ref_bbox
ref_points_detach = inter_ref_bbox.detach(
) if self.training else inter_ref_bbox
return torch.stack(dec_out_bboxes), torch.stack(dec_out_logits)
其中,核心代码如下:
单层DecoderLayer(TransformerDecoderLayer)
在该部分中,数据输入单层DecoderLayer后执行的操作如下:
可变形注意力计算模型(MSDeformableAttention)
可变形注意力模块构造如下:
MSDeformableAttention(
(sampling_offsets): Linear(in_features=256, out_features=192, bias=True)
(attention_weights): Linear(in_features=256, out_features=96, bias=True)
(value_proj): Linear(in_features=256, out_features=256, bias=True)
(output_proj): Linear(in_features=256, out_features=256, bias=True)
)
我们对照着Deformable-DETR
的结构图来观察一下输入参数,首先是Query Feature
,其对应的参数是self.with_pos_embed(tgt, query_pos_embed)
,Reference Point
的维度为torch.Size([4, 498, 1, 4])
,在计算时,我们只选用中心点坐标即可,Input Feature Maps
对应的是memory
,即Encoder输出的特征图。
关于这个过程的代码,我就不在此一一赘述了,我们只需知道最终得到的结果即可。
最终得到可变形交叉注意力的计算结果如下:
单层DecoderLayer结果
将可变形自注意力计算结果拿到后,便是一系列normal等操作,最后返回单层DecoderLayer的结果:
这个结果会进行如下操作:
多层DecoderLayer结果
上述过程是在循环里,代码中有3层,经过多层DecodeLayer计算后,最终得到输出的分类结果与回归结果,将其返回,该部分完整代码如下:
out_bboxes, out_logits = self.decoder(
target,
init_ref_points_unact,
memory,
spatial_shapes,
level_start_index,
self.dec_bbox_head,
self.dec_score_head,
self.query_pos_head,
attn_mask=attn_mask)
最终 Decoder
模块的输出结果如下:
pred_logits:300个特征向量产生的分类结果
pred_boxes:300个特征向量产生的Anchor
aux_outputs,每个Decoder层的结果,因为Decoder中有3层,因此其采用 list 形式存储,每个list中的结果如下:
dn_aux_outputs为每层Decoder加噪查询向量输出结果
此外,还有加噪向量
最终完成了Decoder的计算,接下来便是通过匈牙利匹配方法来匹配预测结果与目标了,同时进行损失计算。
可变形注意力分值计算方法如下:
def deformable_attention_core_func(value, value_spatial_shapes, sampling_locations, attention_weights):
"""
Args:
value (Tensor): [bs, value_length, n_head, c]
value_spatial_shapes (Tensor|List): [n_levels, 2]
value_level_start_index (Tensor|List): [n_levels]
sampling_locations (Tensor): [bs, query_length, n_head, n_levels, n_points, 2]
attention_weights (Tensor): [bs, query_length, n_head, n_levels, n_points]
Returns:
output (Tensor): [bs, Length_{query}, C]
"""
bs, _, n_head, c = value.shape
_, Len_q, _, n_levels, n_points, _ = sampling_locations.shape
split_shape = [h * w for h, w in value_spatial_shapes]
value_list = value.split(split_shape, dim=1)
sampling_grids = 2 * sampling_locations - 1
sampling_value_list = []
for level, (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[level].flatten(2).permute(
0, 2, 1).reshape(bs * n_head, c, h, w)
# N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2
sampling_grid_l_ = sampling_grids[:, :, :, level].permute(
0, 2, 1, 3, 4).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.permute(0, 2, 1, 3, 4).reshape(
bs * n_head, 1, Len_q, n_levels * n_points)
output = (torch.stack(
sampling_value_list, dim=-2).flatten(-2) *
attention_weights).sum(-1).reshape(bs, n_head * c, Len_q)
return output.permute(0, 2, 1)