DETR pytorch代码详解及原理分析

changxl.

已于 2024-08-12 16:58:06 修改

阅读量793

点赞数 20

文章标签： pytorch 人工智能 python

于 2024-08-12 16:44:59 首次发布

本文链接：https://blog.csdn.net/zhuangyuan7838/article/details/141135790

版权

DETR

DETR的过程除了论文，主要借鉴了Bubbliiiing大神的文章及代码，位置在睿智的目标检测65——Pytorch搭建DETR目标检测平台。

DETR，全称为Detection Transformer，是Facebook在ECCV2020上提出的基于Transformer的端到端目标检测网络。最大的特点就是 不需要预定义的先验anchor，也不需要NMS的后处理策略，就可以实现端到端的目标检测。

DETR的总体框架如下，先通过backbone提取图像的特征；再结合位置编码，送入到transformer encoder-decoder中，该编码器解码器的结构基本与transformer相同，主要是在输入部分和输出部分的修改；最后得到类别和bbox的预测，并通过二分匹配计算损失来优化网络。

在这里插入图片描述

整个DETR网络分为四个部分：主干网络、编码器、解码器、预测头。

主干网络backbone：就是常规的目标检测方法进行特征提取，通常进行32倍下采样得到的有效特征层。YOLO经过backbone通常会输出三个有效特征层，DETR或者是其他transformer网络只输出最后一层。将有效特征层结合位置编码（pos-encoding)，作为transformer的输入。
transformer编码器：对于有效特征层进行特征编码，输出编码后的特征层。
transformer解码器：对于编码后的特征层进行特征查询，通过可学习的查询向量加强后的有效特征层进行查询，获得预测结果。
预测头：对预测结果进行维度变换，结果分析。图中画了4个前馈网络（FFN)，实际代码只用了两个。一个用作分类器，判断预测结果的类别。一个用作回归器，对预测结果给出中心点坐标和宽高。

总体推理流程代码如下：

# 传入主干网络中进行预测
# batch_size, 3, 800, 800 => batch_size, 2048, 25, 25
features, pos = self.backbone(samples)

# 将网络的结果进行分割，把特征和mask进行分开
# batch_size, 2048, 25, 25, batch_size, 25, 25
src, mask = features[-1].decompose()
assert mask is not None
# 将主干的结果进行一个映射，然后和查询向量和位置向量传入transformer。
# batch_size, 2048, 25, 25 => batch_size, 256, 25, 25 => 6, batch_size, 100, 256
hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]

# 输出分类信息
# 6, batch_size, 100, 256 => 6, batch_size, 100, 21
outputs_class = self.class_embed(hs)
# 输出回归信息
# 6, batch_size, 100, 256 => 6, batch_size, 100, 4
outputs_coord = self.bbox_embed(hs).sigmoid()
# 只输出transformer最后一层的内容
# batch_size, 100, 21, batch_size, 100, 4
out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}

对结果进行解析时，对于分类结果进行softmax解析，获取分类结果。对于回归结果，由于经过sigmoid调整为0~1区间，直接认为其为归一化后的中心点左边和宽高，调整为左上右下形式，乘以缩放尺度即为最终预测框。

out_logits, out_bbox = outputs['pred_logits'], outputs['pred_boxes']

assert len(out_logits) == len(target_sizes)
assert target_sizes.shape[1] == 2

prob = F.softmax(out_logits, -1)
scores, labels = prob[..., :-1].max(-1)

# convert to [x0, y0, x1, y1] format
boxes = self.box_cxcywh_to_xyxy(out_bbox)
# and from relative [0, 1] to absolute [0, height] coordinates
img_h, img_w    = target_sizes.unbind(1)
img_h           = img_h.float()
img_w           = img_w.float()
scale_fct       = torch.stack([img_w, img_h, img_w, img_h], dim=1)
boxes           = boxes * scale_fct[:, None, :]

outputs = torch.cat([
        torch.unsqueeze(boxes[:, :, 1], -1),
        torch.unsqueeze(boxes[:, :, 0], -1),
        torch.unsqueeze(boxes[:, :, 3], -1),
        torch.unsqueeze(boxes[:, :, 2], -1),
        torch.unsqueeze(scores, -1),
        torch.unsqueeze(labels.float(), -1),    
    ], -1)

results = []
for output in outputs:
    results.append(output[output[:, 4] > confidence])
# results = [{'scores': s, 'labels': l, 'boxes': b} for s, l, b in zip(scores, labels, boxes)]
return results

Transformer的核心就是自注意力机制，以下对DETR的自注意力机制和encoder、decoder进行分析。

自注意力计算流程

生成 Q、 K、 V:
- 计算方式如下：
- Q = K = V
计算注意力权重:
- 使用Q和K计算注意力权重：
- 其中 d_k是键的维度，缩放因子用于稳定计算。
- 通过 Q 和 K 计算注意力权重，然后加权 V 得到输出。

Encoder 层的输入与自注意力机制

输入:
- 特征向量：Encoder 接收的是图像经过 backbone 提取出的特征张量。这些特征张量通常是高维的，包含了图像的各种信息。
- 位置编码:为了使模型理解图像中特征的空间位置信息，通常会在特征张量上添加位置编码（Positional Encoding）。这可以使模型知道特征在原始图像中的位置。

自注意力机制的输入:

查询 Q: 由特征张量和位置编码相加得到。Q = Feature_Tensor + Positional_Encoding
键 *K*: 也是由特征张量和位置编码相加得到。K = Feature_Tensor + Positional_Encoding
值 *V*: 仅为特征张量本身，不加位置编码。V = Feature_Tensor

encoder部分的编码层代码如下：

# 添加位置信息
# 625, batch_size, 256 => 625, batch_size, 256
q = k = self.with_pos_embed(src, pos)
# 使用自注意力机制模块
# 625, batch_size, 256 => 625, batch_size, 256
src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]
# 添加残差结构
# 625, batch_size, 256 => 625, batch_size, 256
src = src + self.dropout1(src2)

# 添加FFN结构
# 625, batch_size, 256 => 625, batch_size, 2048 => 625, batch_size, 256
src = self.norm1(src)
src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
# 添加残差结构
# 625, batch_size, 256 => 625, batch_size, 256
src = src + self.dropout2(src2)
src = self.norm2(src)
return src

Decoder 层的输入与自注意力机制

输入:
- 特征张量：Decoder 接收的是图像经过 backbone 在经过Encoder提取出来的特征张量。
- 位置编码：pos，用于提供位置信息。
- 查询嵌入：query_embedding，通常是解码器的当前状态或目标序列的嵌入。
自注意力机制的输入:
- 查询 Q: query_embedding
- 键 *K*: 由encoder输出的特征张量和位置编码相加得到。
- 值 *V*: 仅为encoder输出的特征张量，不加位置编码。

这样能保证输出的是指定query_embedding个数，且图像经过backbone、encoder、decoder的特征向量，对特征向量进行解码可用于后续做分类和回归。

decoder部分的解码层代码如下：

#---------------------------------------------#
#   q自己做一个self-attention
#---------------------------------------------#
# tgt + query_embed
# 100, batch_size, 256 => 100, batch_size, 256
q = k = self.with_pos_embed(tgt, query_pos)
# q = k = v = 100, batch_size, 256 => 100, batch_size, 256
tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
# 添加残差结构
# 100, batch_size, 256 => 100, batch_size, 256
tgt = tgt + self.dropout1(tgt2)
tgt = self.norm1(tgt)

#---------------------------------------------#
#   q、k、v联合做一个self-attention
#---------------------------------------------#
# q = 100, batch_size, 256, k = 625, batch_size, 256, v = 625, batch_size, 256
# 输出的序列长度以q为准 => 100, batch_size, 256
tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                           key=self.with_pos_embed(memory, pos),
                           value=memory, attn_mask=memory_mask,
                           key_padding_mask=memory_key_padding_mask)[0]
# 添加残差结构
# 100, batch_size, 256 => 100, batch_size, 256
tgt = tgt + self.dropout2(tgt2)
tgt = self.norm2(tgt)
 
#---------------------------------------------#
#   做一个FFN
#---------------------------------------------#
# 100, batch_size, 256 => 100, batch_size, 2048 => 100, batch_size, 256
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
tgt = tgt + self.dropout3(tgt2)
tgt = self.norm3(tgt)
return tgt

在这个过程里会先对q做一个自注意力机制，tgt = torch.zeros_like(query_embed)，我的理解是这个q的自注意力机制是先对query_embed做了个自增强的作用，如果有其他理解请指教。

由于Transformer的注意力机制，网络考虑了整个图像的上下文信息，所以decoder输出的维数和自注意力机制的query_pos强相关。

# 用于传入transformer进行查询的查询向量
self.query_embed    = nn.Embedding(num_queries, hidden_dim)

经过decoder以后输出的也是num_queries个预测框。

这也是DETR区别于传统目标检测网络（例如YOLO）的一个很重要的地方，它输出的直接就是num_queries个预测框，而且结合了上下文信息，所以也不需要做NMS，直接对预测框的得分进行筛选就可以。

那么为什么它输出的这么多预测框就有效呢？

原因是DETR在训练过程中会对预测框和真实框使用匈牙利算法做二分匹配。

匈牙利算法是一种解决二分图最大匹配问题的经典算法，详见https://blog.csdn.net/guoqingru0311/article/details/129142268

python中可直接使用scipy.optimize.linear_sum_assignment完成匈牙利算法的分配。

from scipy.optimize import linear_sum_assignment
import numpy as np
# 代价矩阵
cost =np.array([[0.9,0.6,0,0],[0,0.3,0.9,0],[0.5,0.9,0,0],[0,0,0.2,0]])
# 匹配结果：该⽅法的⽬的是代价最⼩，这⾥是求最⼤匹配，所以将cost取负数
row_ind,col_ind=linear_sum_assignment(-cost)
#对应的⾏索引
print("⾏索引：\n{}".format(row_ind))
#对应⾏索引的最优指派的列索引
print("列索引：\n{}".format(col_ind))
#提取每个⾏索引的最优指派列索引所在的元素，形成数组
print("匹配度：\n{}".format(cost[row_ind,col_ind]))

匈牙利算法是根据代价矩阵，得到线性的匹配的结果。
在这里插入图片描述
如上图中，左侧的0与右侧的1匹配，左侧的1与右侧的2匹配，左侧的2与右侧的1匹配，左侧的3与右侧的3匹配。此时完成匈牙利算法。

所以在DETR中，如何构建代价矩阵，是理解DETR匹配的关键。

由上述解码层可知，解码层输出的维度是[num_queries, hidden_dim]，decoder是由多个解码层组成的（代码中num_decoder_layers=6）。堆叠上batch_size后可以得到Decoder模块输出的维度是[num_decoder_layers, batch_size, num_queries, hidden_dim]。

然后对transformer的结果进行分类和回归

# 输出分类信息
self.class_embed    = nn.Linear(hidden_dim, num_classes + 1)
# 输出回归信息
self.bbox_embed     = MLP(hidden_dim, hidden_dim, 4, 3)

class MLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

可以看出，分类和回归主要都是进行线性变换，分类器输出的是[num_decoder_layers, batch_size, num_queries,num_classes + 1]（+1的意思是多了个背景类别），回归器的输出是[num_decoder_layers, batch_size, num_queries, 4]，4表示中心点坐标及宽高。

只输出Transformer的最后一层（到这个部分感觉之前输出那么多都有点多余了，经过transformer直接获取最后一层解码器的输出就可以了吖，不理解为何要输出num_decoder_layers这个维度）：

# 输出分类信息
outputs_class = self.class_embed(hs)
# 输出回归信息
outputs_coord = self.bbox_embed(hs).sigmoid()
out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}

至此，out中的pred_logits的维度是[ batch_size, num_queries,num_classes + 1]，pred_boxes的维度是[batch_size, num_queries, 4]。

对分类预测结果进行平铺，变换维度为[ batch_size*num_queries,num_classes + 1]，并做softmax。

out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)

在这步就认为softmax后的结果就是该框的分类得分了。

对该batch下所有真实框进行堆叠，获取该batch下的所有真实框，个数为num_target。

tgt_ids = torch.cat([v["labels"] for v in targets])

接下来就可以计算分类成本啦：

# 计算分类成本。预测越准值越小。
cost_class = -out_prob[:, tgt_ids]

也就是取出这batch_size×num_queries数量的框中分类正确的分类得分。加个负号是为了分类得分越高代价越小。cost_class的维度是[batch_size*num_queries, num_target]。

（这个过程是对该batch下的所有预测框对于所有真实框做匹配，如果想要准确的话，应该是每个图像的所有预测框和每个图像的所有真实框做匹配啊，而不应该是该batch下的所有框都做匹配吧？别急，下面有解释。）

接下来对回归结果进行平铺，变换维度为[ batch_size*num_queries,4]。对真实框也做堆叠，变成[num_target, 4]。

# [batch_size * num_queries, 4]
out_bbox = outputs["pred_boxes"].flatten(0, 1)  
# 对真实框进行堆叠
tgt_bbox = torch.cat([v["boxes"] for v in targets])

由于之前输出回归信息时对回归结果进行了sigmoid，也就是变化范围到了0~1之间，所以在创建detaset时对真实框也要做归一化变换，YOLO也会这么做。

接下来计算预测框和真实框的L1成本及IOU成本

# 计算预测框和真实框之间的L1成本。预测越准值越小。
cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)

# 计算预测框和真实框之间的IOU成本。预测越准值越小。
cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))

IOU成本中加上负号也是为了IOU越大的成本越小。计算后的cost_bbox和cost_giou的维度也是[batch_size*num_queries, num_target]。

至此，代价矩阵的三个重要元素都计算完成，且维度都是[batch_size*num_queries, num_target]，也就是该batch下所有预测框对于所有真实框的代价矩阵。分别是：

分类成本
预测框和真实框的L1成本
预测框和真实框的IOU成本

将各项乘以权重，构成最终成本矩阵。

# 最终的成本矩阵
C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou

将成本矩阵转换维度，变为[batch_size，num_queries, num_target]。

C = C.view(bs, num_queries, -1).cpu()

在获取该batch下每个图像中真实框的个数：

sizes = [len(v["boxes"]) for v in targets]

在使用匈牙利算法，完成匹配。

# 对每一张图片进行指派任务，也就是找到真实框对应的num_queries里面最接近的预测结果，也就是指派num_queries里面一个预测框去预测某一个真实框
indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]

这块代码写的特别巧妙，之前已经将成本矩阵转换成了[batch_size，num_queries, num_target]。首先对在C的最后一个维度也就是真实框的个数那个维度用sizes进行split，这样就可以将成本矩阵根据输入图像进行分割。然后c[i]就是每个batch中对应该图像的预测框的成本矩阵，由于sizes的个数肯定是batch，迭代的其实是batch_size这个维度。

至此也就解释了之前我提出来的计算成本矩阵中为什么要批量的计算batch_size×num_queries这个维度，其实在这里是做了分割的。

接下来，在对每张照片使用匈牙利算法linear_sum_assignment进行匹配，获取与各图像真实框最匹配的预测框的索引。

损失计算

有了预测框中与真实框最匹配的索引。那么就可以认为该索引所代表的预测匡是正样本，其他的所有预测框都是负样本，对预测框和真实框计算损失。

通过分类器结果计算分类损失和匹配损失，通过回归器结果计算预测框和真实框的L1-Loss和GIoULoss。

分类损失：使用CELoss进行计算，将匹配框当做正样本，其他框当做负样本。
回归损失：L1-loss和GIoU损失不在赘述。
匹配损失：

  # 计算每个batch真实框的数量
  tgt_lengths     = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
  # 计算不是背景的预测数
  card_pred       = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
  # 然后将不是背景的预测数和真实情况做一个l1损失
  card_err        = F.l1_loss(card_pred.float(), tgt_lengths.float())

核心在于计算不是背景预测数的个数。如果预测框[batch_size, num_queries,num_classes + 1]最大值的索引不是背景，那么判定为预测出前景了，即为判断输出为正样本。用每个图片预测出前景的数量和真实框都是数量做个L1-Loss。

但是实际代码中会计算的损失权重就是{‘loss_ce’: 1, ‘loss_bbox’: 5, ‘loss_giou’: 2}这三个，匹配损失没有被计算。

推理过程

DETR的推理过程在最开头有所提及，较YOLO更简单快捷。前处理部分还是转换维度，做size调整，归一化等工作。然后送入模型推理。推理后得到分类和回归结果，维度分别是[1， num_queries, num_classes + 1]和[1, num_queries, 4]。

然后对预测框分类结果进行softmax获取分类得分，根据预测框分类得分的最大值及索引获取该框的得分和类别。

prob = F.softmax(out_logits, -1)
scores, labels = prob[..., :-1].max(-1)

对预测框的回归结果还原到原图，即从归一化中心点坐标及宽高还原成左上角右下角坐标。

# convert to [x0, y0, x1, y1] format
        boxes = self.box_cxcywh_to_xyxy(out_bbox)
        # and from relative [0, 1] to absolute [0, height] coordinates
        img_h, img_w    = target_sizes.unbind(1)
        img_h           = img_h.float()
        img_w           = img_w.float()
        scale_fct       = torch.stack([img_w, img_h, img_w, img_h], dim=1)
        boxes           = boxes * scale_fct[:, None, :]

        outputs = torch.cat([
                torch.unsqueeze(boxes[:, :, 1], -1),
                torch.unsqueeze(boxes[:, :, 0], -1),
                torch.unsqueeze(boxes[:, :, 3], -1),
                torch.unsqueeze(boxes[:, :, 2], -1),
                torch.unsqueeze(scores, -1),
                torch.unsqueeze(labels.float(), -1),    
            ], -1)

在对得分根据阈值进行筛选，即为推理结果。

results = []
        for output in outputs:
            results.append(output[output[:, 4] > confidence])

DETR的推理过程到此就结束了，把经过softmax后的分类值当做整个预测框的得分，直接做筛选即可，不需要在做NMS。

DETR目标检测总结

DETR 将目标检测视为集合预测问题，使用一种匹配方法（如匈牙利算法）来将预测的框与真实框进行匹配。模型会生成固定数量的预测框，并通过匈牙利算法来确保预测框与真实边界框的匹配，这种方式自然地解决了重叠和冗余问题。DETR 的设计使其能够在无须 NMS 的情况下直接生成高质量的目标检测结果，通过全局上下文、集合预测和有效的匹配算法，减少了重叠和冗余框的问题。这种简化的处理方式是其相较于传统目标检测方法的一个主要优势。

DETR 通过 Transformer 框架直接从输入图像中预测边界框和相应类别。它不仅输出边界框的位置，还同时预测每个边界框的类别信息。这使得模型能够在生成候选框的同时处理重叠框的问题。DETR 输出的每个边界框具有独特的预测分数，通过这种方式，模型能有效地处理多个目标，即使它们重叠。此外，由于模型的设计，多个框可能会得到不同的类别预测，从而进一步减少了重叠框的问题。

DETR 使用自注意力机制，这意味着模型能够全局地考虑图像中所有的区域。在考虑整个图像的上下文信息后，模型能够更有效地判断哪些边界框是冗余的、哪些是最相关的。

DETR 虽然使用 Softmax，而不是 Sigmoid，但其设计考虑了全局上下文和匹配机制（主要还是依托transformer的全局上下文信息特性），能够在无须 NMS 的情况下获得高质量的检测结果。通过匹配机制，DETR 自然过滤掉了低置信度的预测框，从而保持了高效的检测质量。