• models
  • position_encoding.py–存储三角函数实现的位置编码和二维空间位置编码在绝对位置学习
  • transformer.py–区别与原版transformer，主要是在encoder和decedr中的q、k、v的设置方式
  • matcher.py–使用匈牙利算法从输出的100个集合中给每个ground truth匹配合适的目标框
  • backbone.py–使用torch内置的resnet作为backbone抽取图像的feature map
  • detr.py–
  • segmentation.py–

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved """ Various positional encodings for the transformer. """ import math import torch from torch import nn  from util.misc import NestedTensor   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. 类似于transformer三角函数实现的标准位置嵌入层，将每个位置的每个维度映射到角度，即scale默认为0的参数~2Π """     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  # 角度范围为0~2Π         self.scale = scale      def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors  # 获取图像中的图像张量，（batch_size，c，h，w）
        mask = tensor_list.mask  # mask为True表示该部分对应的像素值是通过padding构造的，（batch_size，h，w）
        assert mask is not None
        not_mask = ~mask  # 此处取反，即not_mask为True对应的像素值就是实际图像的像素值
        y_embed = not_mask.cumsum(1, dtype=torch.float32)  # not_mask在列（垂直）方向累加，数据由布尔型转为浮点型，（batch_size，h，w）
        x_embed = not_mask.cumsum(2, dtype=torch.float32)  # not_mask在行（水平）方向累加，数据由布尔型转为浮点型，（batch_size，h，w）
        # 根据上面的行列累加操作，每行每列的值映射为不同的值
        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  # 行方向做归一化

        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  # （batch_size，h，w，num_pos_feats）
        pos_y = y_embed[:, :, :, None] / dim_t  # （batch_size，h，w，num_pos_feats）
        # 在最后一维中，偶数维上使用正弦函数，奇数维上使用余弦函数
        # （batch_size，h，w，num_pos_feats//2）+（batch_size，h，w，num_pos_feats//2）->（batch_size，h，w，num_pos_feats//2，2）->（batch_size，h，w，2*（num_pos_feats//2））
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        # （batch_size，h，w，num_pos_feats//2）+（batch_size，h，w，num_pos_feats//2）->（batch_size，h，w，num_pos_feats//2，2）->（batch_size，h，w，2*（num_pos_feats//2））
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        # （batch_size，h，w，2 *（num_pos_feats // 2））+（batch_size，h，w，2 *（num_pos_feats // 2））->（batch_size，h，w，2*num_pos_feats）->（batch_size，2*num_pos_feats，h，w）
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos


class PositionEmbeddingLearned(nn.Module):
    """ Absolute pos embedding, learned.可以学习的绝对位置嵌入层 """
    def __init__(self, num_pos_feats=256):
        super().__init__()
        # 行列编码的第一个维度都是50表示默认backbone采集的图像feature map不大于50X50
        self.row_embed = nn.Embedding(50, num_pos_feats)
        self.col_embed = nn.Embedding(50, num_pos_feats)
        self.reset_parameters()

    def reset_parameters(self):  # 设置嵌入层参数权重
        nn.init.uniform_(self.row_embed.weight)
        nn.init.uniform_(self.col_embed.weight)

    def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors
        h, w = x.shape[-2:]
        # 一行中的每个位置
        i = torch.arange(w, device=x.device)  # 因为是对行进行编码，就是对该行上所有的列值进行编码
        # 一列中的每个位置
        j = torch.arange(h, device=x.device)  # 因为是对列进行编码，就是对该列上所有的行值进行编码
        x_emb = self.col_embed(i)  # 将一行中的每个列位置进行编码，（w， num_pos_feats）
        y_emb = self.row_embed(j)  # 将一列中的每个行位置进行编码，（h， num_pos_feats）
        pos = torch.cat([
            # （w， num_pos_feats）-> （1， w， num_pos_feats）-> （h， w， num_pos_feats）；即将x_emb按行（垂直）方向重复h此，所有行同一列的编码结果一样
            x_emb.unsqueeze(0).repeat(h, 1, 1),
            # （h， num_pos_feats）-> （h， 1， num_pos_feats）-> （h， w， num_pos_feats）；即将y_emb按列（水平）方向重复h此，所有列同一行的编码结果一样
            y_emb.unsqueeze(1).repeat(1, w, 1), ],
            dim=-1  # 将上面两个向量在最后一维上concat，即（h， w， num_pos_feats）+（h， w， num_pos_feats）->（h， w， 2*num_pos_feats）
        ).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)  # （h， w， 2*num_pos_feats）->（2*num_pos_feats，h， w）->（1，2*num_pos_feats，h， w）->（batch_size，2*num_pos_feats，h， w）
        return pos


def build_position_encoding(args):
    N_steps = args.hidden_dim // 2  # 因为对行、列分类编码，所以维度是隐状态维度的一半
    if args.position_embedding in ('v2', 'sine'):  # 构建三角函数实现的路径位置编码
        # TODO find a better way of exposing other arguments
        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
    elif args.position_embedding in ('v3', 'learned'):  # 构建可学习的绝对路径二维空间位置编码
        position_embedding = PositionEmbeddingLearned(N_steps)
    else:
        raise ValueError(f"not supported { 
          args.position_embedding}")

    return position_embedding

其中三角函数位置编码的计算公式如下图所示

DETR的encoder和decoder结构如下图所示；其中spatial positional encoding是二维空间位置编码方法，该编码分别被加入到encoder的self attention和decoder的cross attention，同时object queires也被加入到decoder的两个attention中