【DETR源码解析】一、整体模型解析

前言

最近在看DETR的源码，断断续续看了一星期左右，把主要的模型代码理清了。一直在考虑以什么样的形式写一写DETR的源码解析。考虑的一种形式是像之前写的YOLOv5那样的按文件逐行写，一种是想把源码按功能模块串起来。考虑了很久还是决定按第二种方式，一是因为这种方式可能会更省时间，另外就是也方便我整体再理解一下吧。

我觉得看代码就是要看到能把整个模型分功能拆开，最后再把所有模块串起来，这样才能达到事半功倍。

另外一点我觉得很重要的是：拿到一个开源项目代码，要有马上配置环境能够正常运行Debug，并且通过解析train.py马上找到主要模型相关的内容，然后着重关注模型方面的解析，像一些日志、计算mAP、画图等等代码，完全可以不看，可以省很多时间，所以以后我讲解源码都会把无关的代码完全剥离，不再讲解，全部精力关注模型、改进、损失等内容。

主要涉及models/detr.py。

Github注释版源码：HuKai97/detr-annotations

一、DETR整体架构

整个搭建过程分为：

1. 搭建DETR：Backbone + Transformer + MLP
2. 初始化损失函数：criterion + 初始化后处理：postprocessors

def build(args):
    # the `num_classes` naming here is somewhat misleading.
    # it indeed corresponds to `max_obj_id + 1`, where max_obj_id
    # is the maximum id for a class in your dataset. For example,
    # COCO has a max_obj_id of 90, so we pass `num_classes` to be 91.
    # As another example, for a dataset that has a single class with id 1,
    # you should pass `num_classes` to be 2 (max_obj_id + 1).
    # For more details on this, check the following discussion
    # https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        # for panoptic, we just add a num_classes that is large enough to hold
        # max_obj_id + 1, but the exact value doesn't really matter
        num_classes = 250
    device = torch.device(args.device)

    # 搭建backbone resnet + PositionEmbeddingSine
    backbone = build_backbone(args)

    # 搭建transformer
    transformer = build_transformer(args)

    # 搭建整个DETR模型
    model = DETR(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=args.aux_loss,
    )

    # 是否需要额外的分割任务
    if args.masks:
        model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))

    # HungarianMatcher()  二分图匹配
    matcher = build_matcher(args)

    # 损失权重
    weight_dict = {'loss_ce': 1, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    if args.masks:   # 分割任务  False
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
    # TODO this is a hack
    if args.aux_loss:   # 辅助损失  每个decoder都参与计算损失  True
        aux_weight_dict = {}
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
        weight_dict.update(aux_weight_dict)

    losses = ['labels', 'boxes', 'cardinality']
    if args.masks:
        losses += ["masks"]

    # 定义损失函数
    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
                             eos_coef=args.eos_coef, losses=losses)
    criterion.to(device)

    # 定义后处理
    postprocessors = {'bbox': PostProcess()}

    # 分割
    if args.masks:
        postprocessors['segm'] = PostProcessSegm()
        if args.dataset_file == "coco_panoptic":
            is_thing_map = {i: i <= 90 for i in range(201)}
            postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors

二、搭建DETR

class DETR(nn.Module):
    """ This is the DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See transformer.py
            num_classes: number of object classes
            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
                         DETR can detect in a single image. For COCO, we recommend 100 queries.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
        """
        super().__init__()
        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        # 分类
        self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        # 回归
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        # self.query_embed 类似于传统目标检测里面的anchor 这里设置了100个  [100,256]
        # nn.Embedding 等价于 nn.Parameter
        self.query_embed = nn.Embedding(num_queries, hidden_dim)
        self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)
        self.backbone = backbone
        self.aux_loss = aux_loss   # True

    def forward(self, samples: NestedTensor):
        """ The forward expects a NestedTensor, which consists of:
               - samples.tensor: batched images, of shape [batch_size x 3 x H x W]
               - samples.mask: a binary mask of shape [batch_size x H x W], containing 1 on padded pixels

            It returns a dict with the following elements:
               - "pred_logits": the classification logits (including no-object) for all queries.
                                Shape= [batch_size x num_queries x (num_classes + 1)]
               - "pred_boxes": The normalized boxes coordinates for all queries, represented as
                               (center_x, center_y, height, width). These values are normalized in [0, 1],
                               relative to the size of each individual image (disregarding possible padding).
                               See PostProcess for information on how to retrieve the unnormalized bounding box.
               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
                                dictionnaries containing the two above keys for each decoder layer.
        """
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        # out: list{0: tensor=[bs,2048,19,26] + mask=[bs,19,26]}  经过backbone resnet50 block5输出的结果
        # pos: list{0: [bs,256,19,26]}  位置编码
        features, pos = self.backbone(samples)

        # src: Tensor [bs,2048,19,26]
        # mask: Tensor [bs,19,26]
        src, mask = features[-1].decompose()
        assert mask is not None

        # 数据输入transformer进行前向传播
        # self.input_proj(src) [bs,2048,19,26]->[bs,256,19,26]
        # mask: False的区域是不需要进行注意力计算的
        # self.query_embed.weight  类似于传统目标检测里面的anchor 这里设置了100个
        # pos[-1]  位置编码  [bs, 256, 19, 26]
        # hs: [6, bs, 100, 256]
        hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]

        # 分类 [6个decoder, bs, 100, 256] -> [6, bs, 100, 92(类别)]
        outputs_class = self.class_embed(hs)
        # 回归 [6个decoder, bs, 100, 256] -> [6, bs, 100, 4]
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        if self.aux_loss:   # True
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        # dict: 3
        # 0 pred_logits 分类头输出[bs, 100, 92(类别数)]
        # 1 pred_boxes 回归头输出[bs, 100, 4]
        # 3 aux_outputs list: 5  前5个decoder层输出 5个pred_logits[bs, 100, 92(类别数)] 和 5个pred_boxes[bs, 100, 4]
        return out

    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

详细源码解析： 【DETR源码解析】二、Backbone模块 和 【DETR源码解析】三、Transformer模块

三、损失函数 + 后处理

详细源码解析： 【DETR源码解析】四、损失计算和后处理模块

四、源码学习重点

1. backbone：Positional Encoding（PositionEmbeddingSine）；
2. Transformer：TransformerEncoderLayer + TransformerDecoderLayer；
3. 损失函数：匈牙利算法，二分图匹配（self.matcher）
4. 后处理：PostProcess

Reference

官方源码: https://github.com/facebookresearch/detr

b站源码讲解: 铁打的流水线工人

知乎【布尔佛洛哥哥】: DETR 源码解读

CSDN【在努力的松鼠】源码讲解: DETR源码笔记（一）

CSDN【在努力的松鼠】源码讲解: DETR源码笔记（二）

CSDN: Transformer中的position encoding(位置编码一)

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（一）、概述与模型推断】

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（二）、模型训练过程与数据处理】

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（三）、Backbone与位置编码】

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（四）、Detection with Transformer】

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（五）、loss函数与匈牙利匹配算法】

知乎CV不会灰飞烟灭-【源码解析目标检测的跨界之星DETR（六）、模型输出与预测生成】