DETR 模型结构源码

目录

• DETR 模型结构源码
  • End-to-End Object Detection with Transformers（DETR）
  • 模型整体结构
    • 模型构建
  • backbone
  • transformer
    • transformer整体构建
    • Encoder
      • TransformerEncoder类
      • TransformerEncoderLayer类
    • Decoder
      • TransformerDecoder类
      • TransformerDecoderLayer类
  • FFN

End-to-End Object Detection with Transformers（DETR）

论文地址：https://arxiv.org/abs/2005.12872

源代码位置： https://github.com/facebookresearch/detr

参考文献： https://www.cnblogs.com/Glucklichste/p/14057005.html

模型整体结构

论文中模型结构

主干网络

• backbone(CNN-Resnet)
  • CNN网络
  • positional(位置信息)
• transformer
  • encoder
  • decoder
• predicttion head

模型构建

models/detr.py
#  构建两大模型
#  backbone = build_backbone(args)
#  transformer = build_transformer(args)
#  模型连接  DETR
#

def build(args):

    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 和 构建 transformer 
    backbone = build_backbone(args)
    transformer = build_transformer(args)

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

backbone

cnn骨架特征提取
backbone的输入和输出

• input shape=(N,3,W,H)
• output shape=(N,2048,W/32,H/32) #针对 Resnet50 C=2048, 针对 Resnet18,Resnet34 C=512

假设输入是(N,C,H,W)，则resnet50输出是(N,2048,H//32,W//32)，1024比较大，
为了节省计算量，先采用1x1卷积降维为256,(hidden_dim=256,在main.py 中设置参数)
最后转化为序列格式输入到transformer中，输入shape=(H*W,N,256)，H=H/32,W=W/32

class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool,
                 dilation: bool):
        backbone = getattr(torchvision.models, name)(
            replace_stride_with_dilation=[False, False, dilation],
            pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)

        # 针对不同的网络，选择了不同的输出大小
        num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)

···


```python
在 DETR 类中
src 为 backone 的输出 shape=(N,512,W/32,H/32)
# self.input_proj(src) 将 shape=(N,512,W/32,H/32) -> shape=(N,256,W/32,H/32)

hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]

位置信息标注,包含了x,y两个方向的位置信息。编码方式任然采用sincos, 语音序列只是包含了一个方向的位置信息
PositionEmbeddingSine.forward的输入和输出

• input NestedTensor型数据 tensor_list的类型是NestedTensor，内部自动附加了mask，
• x.tensors.shape=((N, 512,W/32, H/32) x.mask.shape=(N,W/32,H/32)
• output: pos.shape=(N, 256, W/32,H/32)

class PositionEmbeddingSine(nn.Module):
    """
    This is a more standard version of the position embedding, very similar to the one
    used by the Attention is all you need paper, generalized to work on images.
    """
    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
        super().__init__()
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

    def forward(self, tensor_list: NestedTensor):

        x = tensor_list.tensors
        mask = tensor_list.mask
        #x.tensors.shape=((N, 512,W/32, H/32)   x.mask.shape=(N,W/32,H/32)

        assert mask is not None
        not_mask = ~mask
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
        
        # 前面输入向量是256，编码是一半sin，一半cos
        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)

        # pos.shape=(N, 256, W/32,H/32)  前128是y方向编码，而128是x方向编码

        return pos

transformer

transformer整体构建

model/transformer.py
Transformer  模型构建
包含 encoder   decoder

class Transformer(nn.Module):

    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()

        # 编码
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)

        # 解码
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec)

        self._reset_parameters()

        self.d_model = d_model
        self.nhead = nhead

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, src, mask, query_embed, pos_embed):
        # flatten NxCxHxW to HWxNxC
		# inputs:  {src,mask,query_embed,pos} 由 DETR.forward 获取来自 backbone
        
        bs, c, h, w = src.shape
		
        # 先对数据做变换
        # 特殊说明 这里是经过backbone 输出的特征 (N,256,W/32,H/32) 之后transformer过程中 输出shape为(H/32xW/32,N,256) 特征的宽和高没有变化，为了书写方法方便，我这里将 W/32,H/32 写成为 W,H 
        # src=(N,256,W/32,H/32)-> (WH,N,256)
        # pos_embed=(N,256,W,H)-> (WH,N,256)
        # query_embed=(100,256) -> (100,N,256)
        # mask=(N,W,H) -> (N,WH)
        src = src.flatten(2).permute(2, 0, 1)
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
        mask = mask.flatten(1)
		
        # 解码  第一层 首次参数设置为0，后续自动更新
        tgt = torch.zeros_like(query_embed)

        # encoder  src=(WH,N,256) mask= (N,WH)  pos_embed= (WH,N,256)
        # 输出 (WH,N,256)
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
       

        # decoder tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) 
        # pos_embed=(WH,N,256) query_embed=(100,N,256)
        # 输出 hs=(decoder_layers, 100, N, 256)
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)

        
        # return (decoder_layers, N, 100, 256) (N, 256, H, W])
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)

Encoder

编码器结构和输入输出
编码器的输入有三个 src=(WH,N,256) src_mask= (N,WH) pos_embed= (WH,N,256) 注释：W=W/32,H=H/32

• 由图像生成的序列,shape=(WH,N,256)
• 掩码信息,shape= (N,WH)
• 图像序列的空间位置信息,shape=(WH,N,256)

经过6层编码后 输出只有一个 序列,shape和输入的src 序列保持一直,shape=(WH,N,256) 注释：W=W/32,H=H/32

模型细节

• 原始transformer的n个编码器输入中，只有第一个编码器需要输入位置编码向量，但是DETR里面对每个编码器都输入了同一个位置编码向量
• QKV处理逻辑不同，在编码器内部位置编码仅仅和 Q K 相加，V 不做任何处理

TransformerEncoder类

def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

class TransformerEncoder(nn.Module):

    def __init__(self, encoder_layer, num_layers, norm=None):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm

    def forward(self, src,
                mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        output = src
        
        #  默认设置了 6个 编码器，循环6遍
        #  encoder input  src=(WH,N,256) src_mask= (N,WH)  pos_embed= (WH,N,256)
        #  output -> output  (WH,N,256)
        # 包含了多层相同的结构,首尾相连,上一层输出为下一层的输入
        
        for layer in self.layers:
            output = layer(output, src_mask=mask,
                           src_key_padding_mask=src_key_padding_mask, pos=pos)
        if self.norm is not None:
            output = self.norm(output)

        return output

TransformerEncoderLayer类

class TransformerEncoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self,
                     src,
                     src_mask: Optional[Tensor] = None,
                     src_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None):
        
        # src=(WH,N,256) mask= (N,WH)  pos_embed= (WH,N,256)
        #  with_pos_embed  输入是 src  pos  {图片序列,位置信息}
        # 对 Q K 进行更新
        q = k = self.with_pos_embed(src, pos)

        # MultiheadAttention 多头注意力机制
        # 在编码器内部位置编码仅仅和QK相加，V不做任何处理

        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]
        # 残差
        src = src + self.dropout1(src2)
        src = self.norm1(src)
        # FFN
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src

    def forward_pre(self, src,
                    src_mask: Optional[Tensor] = None,
                    src_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None):
        src2 = self.norm1(src)
        q = k = self.with_pos_embed(src2, pos)
        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
                              key_padding_mask=src_key_padding_mask)[0]

        src = src + self.dropout1(src2)
        src2 = self.norm2(src)
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
        src = src + self.dropout2(src2)

        return src

    def forward(self, src,
                src_mask: Optional[Tensor] = None,
                src_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None):
        # encoder  src=(WH,N,256) mask= (N,WH)  pos_embed= (WH,N,256)
        # output=(WH,N,256)
        #  默认 normalize_before=False 只对 forward_post 函数注解
        if self.normalize_before:
            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)

        return self.forward_post(src, src_mask, src_key_padding_mask, pos)

Decoder

解码器结构和输入输出

输入参数
解码器的输入 有五个参数 decoder tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) pos_embed=(WH,N,256) query_pos=(100,N,256)

• tgt 可以理解为上一层解码器的解码输出 shape=(100,N,256) 第一层的tgt=torch.zeros_like(query_embed) 为零矩阵
• memory 最后一个编码器输出 shape=(WH,N,256)
• mask 掩码信息 shape=(N,WH)
• pos 和编码器输入中完全相同位置参数 shape=(WH,N,256)
• query_pos 是可学习输出位置向量， 个人理解 解码器中的这个参数 全局共享 提供全局注意力 query_pos=(100,N,256)

输出参数

• 输出 (decoder_layers, 100, N, 256) decoder_layers 为解码器的数量(层数)，原文默认设置为6层

原始transformer顺序解码操作不同的是，detr一次就把N个无序框并行输出

Obeject Query
针对 query_pos 参数的其他博客解释
论文中指出object queries作用非常类似faster rcnn中的anchor，只不过这里是可学习的，不是提前设置好的。
object queries(shape是(100,256)) 源代码中，这是一个torch.nn.Embedding的对象。
官方介绍：一个保存了固定字典和大小的简单查找表。这个模块常用来保存词嵌入和用下标检索它们。模块的输入是一个下标的列表，输出是对应的词嵌入。

个人理解：query_pos 可以简单认为是输出位置编码,其作用主要是在学习过程中提供目标对象和全局图像之间的关系,相当于全局注意力，必不可少非常关键。代码形式上是可学习位置编码矩阵。和编码器一样，该可学习位置编码向量也会输入到每一个解码器中。我们可以尝试通俗理解：object queries矩阵内部通过学习建模了100个物体之间的全局关系，并且参与到网络的学习当中。

其他细节:

• tgt(第一次输入是query embeding，第二次是上一层的输出out);
• 和编码器一样，只是Q 与 K加上了位置编码信息, V不会加入位置编码
• 引入可学习的Object queries
• 不需要顺序解码，一次即可输出N个无序集合

TransformerDecoder类

class TransformerDecoder(nn.Module):

    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
        super().__init__()
        self.layers = _get_clones(decoder_layer, num_layers)
        self.num_layers = num_layers
        self.norm = norm
        self.return_intermediate = return_intermediate

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):
        # decoder tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) pos_embed=(WH,N,256) query_embed=(100,N,256)
        output = tgt
        intermediate = []

        for layer in self.layers:

            output = layer(output, memory, tgt_mask=tgt_mask,
                           memory_mask=memory_mask,
                           tgt_key_padding_mask=tgt_key_padding_mask,
                           memory_key_padding_mask=memory_key_padding_mask,
                           pos=pos, query_pos=query_pos)
            if self.return_intermediate:
                intermediate.append(self.norm(output))

        if self.norm is not None:
            output = self.norm(output)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)

        # intermediate=[outpout...]     intermediate[0].shape=(100,N,256)
        # return_intermediate = True
        if self.return_intermediate:
            return torch.stack(intermediate)

        return output.unsqueeze(0)

TransformerDecoderLayer类

class TransformerDecoderLayer(nn.Module):

    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

        self.activation = _get_activation_fn(activation)
        self.normalize_before = normalize_before

    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
        return tensor if pos is None else tensor + pos

    def forward_post(self, tgt, memory,
                     tgt_mask: Optional[Tensor] = None,
                     memory_mask: Optional[Tensor] = None,
                     tgt_key_padding_mask: Optional[Tensor] = None,
                     memory_key_padding_mask: Optional[Tensor] = None,
                     pos: Optional[Tensor] = None,
                     query_pos: Optional[Tensor] = None):

        #  # decoder tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) pos_embed=(WH,N,256) query_embed=(100,N,256)
        # 解码 第一次注意力机制  tgt=(100,N,256)  是 上一个单元输出 如果是第一次 torch.zeros_like(query_embed)
        # query_embed=(100,N,256)  query_pos 应该是共享单元，不管多少层都是公用一组数据


        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        # multihead_attn
        # query=self.with_pos_embed(tgt, query_pos)   在第二次注意力机制中 对 Q 进行更新
        # key=self.with_pos_embed(memory, pos)     在第二次注意力机制中对 K 进行更新

        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]
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)
        # FFN
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
        tgt = tgt + self.dropout3(tgt2)
        tgt = self.norm3(tgt)
        return tgt

    def forward_pre(self, tgt, memory,
                    tgt_mask: Optional[Tensor] = None,
                    memory_mask: Optional[Tensor] = None,
                    tgt_key_padding_mask: Optional[Tensor] = None,
                    memory_key_padding_mask: Optional[Tensor] = None,
                    pos: Optional[Tensor] = None,
                    query_pos: Optional[Tensor] = None):
        #  # decoder tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) pos_embed=(WH,N,256) query_embed=(100,N,256)

        tgt2 = self.norm1(tgt)
        q = k = self.with_pos_embed(tgt2, query_pos)
        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
                              key_padding_mask=tgt_key_padding_mask)[0]
        tgt = tgt + self.dropout1(tgt2)
        tgt2 = self.norm2(tgt)
        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
                                   key=self.with_pos_embed(memory, pos),
                                   value=memory, attn_mask=memory_mask,
                                   key_padding_mask=memory_key_padding_mask)[0]
        tgt = tgt + self.dropout2(tgt2)
        tgt2 = self.norm3(tgt)
        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
        tgt = tgt + self.dropout3(tgt2)
        return tgt

    def forward(self, tgt, memory,
                tgt_mask: Optional[Tensor] = None,
                memory_mask: Optional[Tensor] = None,
                tgt_key_padding_mask: Optional[Tensor] = None,
                memory_key_padding_mask: Optional[Tensor] = None,
                pos: Optional[Tensor] = None,
                query_pos: Optional[Tensor] = None):

        # decoder input tgt=(100,N,256) memory=(WH,N,256),mask=(N,WH) pos_embed=(WH,N,256) query_embed=(100,N,256)

        # ISFalse
        if self.normalize_before:
            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)

FFN

最后是接了一个FFN,就是两个全连接层，一个用于分类，一个用于回归预测

分类: 一层模型结构
最终预测 MLP模型 是由具有ReLU激活功能且具有隐藏层的3层感知器和线性层计算的。 FFN预测框的标准化中心坐标，高度和宽度， 输入图像，然后线性层使用softmax函数预测类标签

DETR类中
        # 输入  hs.shape = (decoder_layers, N, 100, 256)

        # 分类 self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
        # FFN  Linear class   input=(decoder_layers, N, 100, 256)         output=(decoder_layers, N, 100, num_classes+1)
        outputs_class = self.class_embed(hs)
        
        # 预测  self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
        # MLP  Bounding box   input=(decoder_layers, N, 100, 256)         output=(decoder_layers, N, 100, 4)
        outputs_coord = self.bbox_embed(hs).sigmoid()