detr源码解读

y0k1n0

已于 2024-07-26 14:54:34 修改

阅读量617

点赞数 11

分类专栏：论文阅读深度学习文章标签： python 深度学习目标检测计算机视觉 pytorch

于 2024-07-25 15:11:04 首次发布

本文链接：https://blog.csdn.net/yosh1n0/article/details/140690791

版权

论文阅读同时被 2 个专栏收录

3 篇文章 0 订阅

订阅专栏

深度学习

3 篇文章 0 订阅

订阅专栏

detr 代码解读

position_encoding

PositionEmbeddingSine

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):  # 指示tensor_list的类型应该为NestedTensor
        x = tensor_list.tensors
        mask = tensor_list.mask  # mask是三维张量
        assert mask is not None
        not_mask = ~mask  # 表示非填充区域
        y_embed = not_mask.cumsum(1, dtype=torch.float32)  # 在每个三维张量的第二维（二维数组的x）累加 [batch, h, w]
        x_embed = not_mask.cumsum(2, dtype=torch.float32)  # 在每个三维张量的第三维（二维数组的y）累加 [batch, 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)  # dim_t // 2 确保相邻的两个结果相同，用于sin和cos
                                                                             # 除以特征数用于归一化操作
                                                                             # 为每个位置编码维度计算一个缩放因子，
                                                                             # 使得位置编码在不同维度上有不同的频率。\
                                                                             # [num_pos_feats]

        pos_x = x_embed[:, :, :, None] / dim_t  # 添加维度方便广播运算 [batch, h, w, num_pos_feats]
        pos_y = y_embed[:, :, :, None] / dim_t  # [batch, h, w, num_pos_feats]
        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)
        return pos

这段代码利用掩码的累加作为位置编码的pos，使用如下公式利用三角函数表达式完成位置编码：
$PE_{(pos,2i)}=sin(\frac{pos}{10000^{\frac{2i}{d_model}}}) \\ PE_{(pos,2i + 1)}=cos(\frac{pos}{10000^{\frac{2i}{d_model}}})$
其中，pos表示token在序列中的位置，设句子长度为 L，则 $pos=0,1,\ldots,L-1$ 是token的位置向量， $PE (p os, 2 i)$ 表示这个位置向量里的第i个元素， $2 i + 1$ 表示奇数维度， $2 i$ 表示偶数维度； $d_{model}$ 表示token的维度。

PositionEmbeddingSine该函数使用的是上述位置编码的二维化版本，利用位置掩码mask（指示是否为填充区域）的x方向和y方向上累加后得到的归一化作为一组正交的位置编码，将他们经过三角函数运算后拼接到一起作为一个具有 $2\times 特征维度$ 的位置编码坐标。

PositionEmbeddingLearned

class PositionEmbeddingLearned(nn.Module):
    """
    自学习的绝对位置编码
    """
    def __init__(self, num_pos_feats=256):
        super().__init__()
        # 需要编码的行列不超过五十个
        self.row_embed = nn.Embedding(50, num_pos_feats)  # nn.Embedding(num_embeddings, embedding_dim) 词表大小， 词向量维度
        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([
            x_emb.unsqueeze(0).repeat(h, 1, 1),  # [h, w, num_pos_feats]
            y_emb.unsqueeze(1).repeat(1, w, 1),  # [h, w, num_pos_feats]
        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)  # [batchsize, 2 * num_pos_feats, h, w]
        return pos

该类可以处理特征小于50x50的图像，通过绝对编码的方式，为每个位置的行（）和列（）分别赋值一个大小为num_pos_feats的位置编码，最后拼接并返回一个大小为[batchsize, 2 * num_pos_feats, h, w]的位置编码。

embedding中的权重矩阵会在训练过程中由反向传播更新。

backbond

整个backbond特征提取实际上是在缩小特征图的大小，因为transformer的时间复杂度是 $O(kn^2)$ ，尽量减少特征图的维数有利于加快后面的处理。

FrozenBatchNorm2d

这段代码定义了一个 forward 方法，属于某个神经网络层类。它对输入张量 x 进行前向传播计算.

class FrozenBatchNorm2d(torch.nn.Module):
    """
    BatchNorm2d,其中批处理统计量和仿射参数是固定的。

    从 torchvision.misc.ops 复制粘贴的，并在 rqsrt 前添加了 eps,
    否则除 torchvision.models.resnet[18,34,50,101] 之外的任何其他模型都会产生 nans。
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))  # register_buffer定义的变量不会注册到parameters，而是直接注册到state_dict
        self.register_buffer("bias", torch.zeros(n))   # 不需要每次训练的时候重新加载
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    def _load_from_state_dict(self, state_dict,  # 状态字典，包含模型的参数和其他状态信息。
                               prefix,  # 前缀字符串，用于构造完整的键名。
                               local_metadata,  # 本地元数据，提供关于模型的一些额外信息。
                               strict,  # 布尔值，如果为 True，表示严格模式，要求状态字典的键完全匹配模型的键。
                               missing_keys,  # 列表，用于存储在状态字典中缺失的键。
                               unexpected_keys,  # 列表，用于存储状态字典中意外出现的键。
                               error_msgs):  # 列表，用于存储加载过程中出现的错误信息。
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]  # 如果该键值存在于状态列表，将他删除

        super(FrozenBatchNorm2d, self)._load_from_state_dict(  # 调用父类的方法
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)  # [1, c, 1, 1] 方便后续广播操作
        b = self.bias.reshape(1, -1, 1, 1)  # [1, c, 1, 1] 方便后续广播操作
        rv = self.running_var.reshape(1, -1, 1, 1)  # [1, c, 1, 1] 保存的运行时方差，用于推理阶段的标准化。
        rm = self.running_mean.reshape(1, -1, 1, 1)  # [1, c, 1, 1] 保存的运行时均值，用于推理阶段的标准化。
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()  # 计算倒数平方根
        bias = b - rm * scale
        return x * scale + bias

这个类继承自 torch.nn.Module，用于实现一种冻结的批归一化层，其中批处理统计量和仿射参数是固定的。初始化方法中self.register_buffer 用于注册固定的参数，这些参数不会被优化器更新，但会在模型保存和加载时存储在 state_dict 中。

前向传播算法中计算缩放因子，公式为
$\text{scale} = \text{weight} \times \frac{1}{\sqrt{\text{running var} + \text{eps}}}$
rsqrt() 表示计算倒数平方根。

计算偏移量，公式为
$\text{bias} = \text{bias} - \text{running mean} \times \text{scale}$
对输入张量 x 进行标准化和缩放偏移操作，并返回结果。

整个前向传播算法使用标准的 Batch Normalization公式计算标准化的结果
$running_mean running_var + e p s × weight + bias \text{output} = \frac{x - \text{running\_mean}}{\sqrt{\text{running\_var} + eps}} \times \text{weight} + \text{bias}$

BackboneBase

class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        for name, parameter in backbone.named_parameters():
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)
        if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)  # 选择返回最后还是包括中间层
        self.num_channels = num_channels

    def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors)  # 返回一个字典
                                             # 键值为str，层的名字
                                             # 值是NestedTensor对象
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None  # 掩码为空程序抛出异常
            mask = F.interpolate(m[None].float(),  # 把掩码变为三维张量[1, h, w]
                                 size=x.shape[-2:]).to(torch.bool)[0]  # 利用插值函数，将掩码大小对其宽高,取[h, w]
            out[name] = NestedTensor(x, mask)
        return out

根据return_interm_layers的真值为字典return_layers赋值，并调用IntermediateLayerGetter函数确定返回参数是否包括中间层。

在前向传播中，将以字典的形式{str, NestedTensor}保存参数传播的最后一层（可能还有中间层）结果，并在随后的循环中，针对保存的每一层利用插值函数更新掩码，使他与特征图的高宽一致，并返回最终的字典{str, NestedTensor}。

Backbone

class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,  # ResNet模型的名称
                 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)  # 调用父类的构造函数

整个特征提取的网络骨架使用的是resnet结构。该类通过参数传递选择具体的网络结构，并在选择resnet18或resnet34时设置通道数量为512，其他时候为2048.

Joiner

class Joiner(nn.Sequential):
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):
        xs = self[0](tensor_list)  # 调用第一个子模块（backbone）得到{名称，特征图}的字典
        out: List[NestedTensor] = []  # 初始化一个列表用于存储输出结果，输出的每个结果都是一个特征图
        pos = []  # 空列表用于存储位置编码
        for name, x in xs.items():  # 遍历网络骨架得到的所有特征图
            out.append(x)
            # position encoding
            pos.append(self[1](x).to(x.tensors.dtype))  # 调用位置编码模块，根据特征图生成位置编码

        return out, pos

该类实现了将backbone模块和position_embedding结合，在前向传播方法中，首先先将输入序列通过backbone得到一个形状为{名称，特征图}的字典。初始化一个列表存储输出结果（每个结果都是一个NestedTensor）和一个位置编码列表；最后遍历特征图，将所有特征图和他的位置编码（将特征向量通过position_embedding对象得到）返回。

transformer

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):  # 可选的张量，表示位置嵌入。
        q = k = self.with_pos_embed(src, pos)  # 将位置编码嵌入的输入参数传递给查询q和键值k
        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)  # 归一化
        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))   #src 依次通过以下层：
                                                                                # 线性层 self.linear1。
                                                                                # 激活函数 self.activation(默认 relu)。
                                                                                # 丢弃层 self.dropout。
                                                                                # 线性层 self.linear2。
        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]  # 计算自注意力，返回输出(对应[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):
        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)

这个类再建立transformer的encoder编码层，每个编码层都是由以下子层组成（以前归一化为例）：

位置编码嵌入层
多头注意力
规范化层（残差连接，减小梯度消失）
逐位前馈网络
1. 线性层
2. 激活函数
3. dropout
4. 线性层
规范化层（残差连接，减小梯度消失）

该类拥有两个版本的前向传播函数，分别实现了后归一化和前归一化两种不同归一化的传播函数，根据normalize_before参数选择归一化方式。

前归一化：有助于针对图像数据调整输入特征的尺度，使模型在处理图像数据时更加高效和精确。
后归一化：有助于模型更好地保留和学习输入数据之间的关系，同时也有利于保持模型训练的稳定性和高效性。

_get_clones

def _get_clones(module, N):  # 创建一个包含N个相同模块的列表
    return nn.ModuleList(  # 存储这些副本，以便在PyTorch的模型中进行管理和使用。
        [copy.deepcopy(module) for i in range(N)  # deepcopy确保每个副本都是独立的。
         ])

这个函数完成的工作是将编码层（解码层）深复制N次，并返回复制后的结果。

TransformerEncoder

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

        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

该类实现了编码器架构，将输入序列通过N层编码层后经过归一化作为最后的输出信号。

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,  # 目标序列的输入张量，形状为 (sequence_length, batch_size, d_model)
                     memory,  # 编码器输出的记忆张量，形状为 (sequence_length, batch_size, d_model)
                     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):  # 目标序列的位置嵌入张量
        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)  # 归一化
        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)  # 归一化处理
        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):
        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):
        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)

该类定义了解码器层，每个解码器层都是由以下子层组成：

首先现在目标序列中添加位置编码，作为解码器的查询和键
将解码器的查询、键、值（目标序列）输入自注意力汇聚层
dropout+残差连接
归一化层
将位置编码嵌入的目标序列作为查询、位置编码嵌入的编码器输出作为键、编码器输出作为值，送入多头注意力模型
dropout+残差连接
归一化处理
前馈神经网络
1. 全连接层
2. 激活函数
3. dropout
4. 全连接层
dropout+残差连接
归一化处理

同时，与编码器相同，解码器也采用了两个版本的前向传播函数，分别实现了后归一化和前归一化两种不同归一化的传播函数，根据normalize_before参数选择归一化方式。

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):
        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)  # 并将最终归一化传给中间层列表

        if self.return_intermediate:
            return torch.stack(intermediate)  # 返回所有中间层的输出，使用torch.stack将它们堆叠成一个张量。

        return output.unsqueeze(0)  # 返回最终output，并在第一个维度上添加一个新的维度

这个类定义了一个解码器。该解码器的前向传播函数可以将目标序列经过N个解码层，并在每个解码层中，根据return_intermediate参数的真值决定是否返回归一化后的每个编码层的输出结果；然后解码器还可以根据norm的真值决定对于最终的输出是否归一化。然后当return_intermediate为真时，将所有的中间层的输出(包括最后一层的输出)堆叠为张量返回，否则返回最后一层的输出，并在第一个维度之前添加一个新的维度，以匹配torch.stack(intermediate)。

Transformer

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) # 创建编码器层
                                                                                       # 根据normalize_before的真值调用不同的前向传播
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None  # 如果normalize_before为真创建归一化层
        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()  # 使用Xavier均匀初始化所有维度大于1的参数

        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
        bs, c, h, w = src.shape  # 输入特征：[BatchSize, Channel, h, w]
        src = src.flatten(2).permute(2, 0, 1)  # [h * w, batchsize, channel]
        pos_embed = pos_embed.flatten(2).permute(2, 0, 1)  # [h * w, batchsize, channel]
        query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)  # [num_queries, d_model] -> [num_queries, batch_size, d_model]
        mask = mask.flatten(1) #[batchsize, h, w] -> [batchsize, h * w]

        tgt = torch.zeros_like(query_embed) # 初始化目标张量，形状与query_embed相同 [num_queries, batchsize, d_model]
        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)  # 编码层处理 [h*w, batchsize, d_model]
        hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                          pos=pos_embed, query_pos=query_embed)  # 解码层处理 [num_queries, batch_size, d_model]
        return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)  # hs转置为[num_queries, batchsize, d_model]
                                                                              # memory重塑为[batchsize, channel(d_model), h, w]

transtormer类实现了transformer编码解码模型

初始化
1. ·在类的初始化中，定义了编码器，解码器以及各自的归一化层
2. 并使用xavier_uniform_函数初始化所有维度大于1的参数。
前向传播
1. 前向传播中，首先先把输入特征转换为形状[h * w, batchsize, d_model]，以适配序列模型
2. 将查询位置编码和输入特征位置编码形状转换为[num_queries, batch_size, d_model]以匹配输入特征。
3. 将掩码展开为[batchsize, h * w]。
4. 初始化目标特征，形状与输入特征相同[h * w, batchsize, d_model],并全部置为0。
5. 将输入特征通过编码器，得到编码层的输出序列，形状为[h*w, batchsize, d_model]。
6. 将编码层输出序列作为多头注意力的查询和键值和目标特征一同输入解码层，得到最终的注意力汇聚后的结果张量。
7. 返回解码器结果张量和编码器的输出张量，并将它们重塑与输入形状匹配。

misc

NestedTensor

class NestedTensor(object):
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors  # [batchsize, channel, h, w]
        self.mask = mask  # [batchsize, h, w]

    def to(self, device):  # 将NestedTensor中的张量和掩码移动到指定的设备
        # type: (Device) -> NestedTensor # noqa
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            assert mask is not None
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):  # 分别返回目的张量和掩码
        return self.tensors, self.mask

    def __repr__(self):  # 返回内部张量的字符串形式
        return str(self.tensors)

该类封装了图像张量[batchsize, channel, h, w]和位置掩码[batchsize, h, w]。提供了分别返回张量和掩码的方法decompose、将张量移动到设备的方法to、返回内部张量字符串形式的方法__repr__。

nested_tensor_from_tensor_list

def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    # TODO make this more general
    if tensor_list[0].ndim == 3: # 检查当前列表中第一个张量是否是三维图像 [channel, h, w]
        if torchvision._is_tracing():  # 如果正在使用 torchvision 的 ONNX 导出功能，
                                       # 则调用 _onnx_nested_tensor_from_tensor_list 函数来处理张量列表，
                                       # 因为 nested_tensor_from_tensor_list 不能很好地导出到 ONNX。
            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # TODO make it support different-sized images
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])  # 计算张量列表中所有图片最大尺寸
        # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
        batch_shape = [len(tensor_list)] + max_size  # [batchsize, channel, h, w]
        b, c, h, w = batch_shape
        dtype = tensor_list[0].dtype  # 获取数据类型
        device = tensor_list[0].device  # 获取设备信息
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)  # 创建一个形状为[b, c, h, w]的零填充张量
                                                                       # 这个张量用于存储填充后的图像数据。
        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)  # 创建一个形状为[b, h, w]的布尔掩码张量
                                                                       # 所有值初始化为 True。
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)  # 复制img的所有元素到pad_img的对应位置：
            m[: img.shape[1], :img.shape[2]] = False  # 更新掩码张量m，将填充图像区域的掩码值设置为False
    else:
        raise ValueError('not supported')
    return NestedTensor(tensor, mask)

该函数用于将张量的列表转换为定义的NestedTensor对象。通过获取张量列表中所有形状中最大的类别（分别在长宽高上），得到返回张量的大小[张量数量(batch), channel, h, w]作为张量和掩码的形状，并从张量的左上开始填充原图像的所有参数。最后调用构造函数返回NestedTensor对象。

_max_by_axis

def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0]
    for sublist in the_list[1:]:
        for index, item in enumerate(sublist):
            maxes[index] = max(maxes[index], item)
    return maxes

该函数遍历所有张量的每一个维度，找到每一个维度上张量的最大值并返回。

detr

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: 用于特征提取的主干网络模块。
            transformer: Transformer 结构的模块。
            num_classes: 对象类别的数量。
            num_queries: DETR 在单张图像中最多可以检测的对象数量。对于 COCO 数据集，推荐使用 100 个查询。
            aux_loss: 是否使用辅助解码损失（即每个解码器层的损失）。
        """
        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)  # 创立线性层，将隐藏维度映射到类别+1（包括背景类）
        self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)  # 多层感知机，用于边界框预测
        self.query_embed = nn.Embedding(num_queries, hidden_dim)  # 创建查询嵌入，查询数量为num_queries, 每个查询的维度为hidden_dim
        self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)  # 创建一个卷积层，将主干网络的输出通道数映射到 
                                                                                       # Transformer 的隐藏维度
        self.backbone = backbone
        self.aux_loss = aux_loss

    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 是一个列表或 torch.Tensor，则将其转换为 NestedTensor 类型。
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)  # 通过骨干网络提取特征()和位置编码

        src, mask = features[-1].decompose()  # 将提取到的特征分为输出特征和掩码
        assert mask is not None
        hs = self.transformer(self.input_proj(src), # 将backbone的输出特征经过Conv2d，将num_channels转换为hidden_dim
                              mask,   # 这个掩码用于告诉 Transformer 哪些位置应该被忽略(置为负无穷)
                              self.query_embed.weight,  # 查询向量用于查询目标图像中的目标
                              pos[-1]  # 获取最后一层的位置信息
                              )[0]  #  返回编码后的查询嵌入

        outputs_class = self.class_embed(hs)  # 将transformer的输出映射到类别上 [保存的解码器层, batchsize, num_queries, num_classes]
        outputs_coord = self.bbox_embed(hs).sigmoid()  # 将transformer的输出映射到多层感知机上，并用sigmod归一化到[0, 1]范围
                                                       # 用于预测锚框的位置 [batchsize, num_query, 4(中心x, 中心y, h, w)]
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}  # 存储最后一层解码器经过上述处理后的输出
        if self.aux_loss:
            # 如果 aux_loss 参数为真，表示需要计算辅助损失
            # 调用 _set_aux_loss 方法来计算辅助损失并将结果添加到字典 out 中。
            # 辅助损失是在每个解码器层计算的损失，有助于训练过程中更好地传播梯度。
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        return out

    @torch.jit.unused  # 这个装饰器用于标记方法 _set_aux_loss 不会被 TorchScript 使用。
    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])]

这个是整个项目最关键的函数，实现了整体的目标检测模型。

初始化
- num_queries：询问个数
- transformer：transformer解码编码架构
- hidden_dim：隐藏层维度
- class_embed：线性层，将输出的隐藏层维度映射到目标个数
- bbox_embed:多层感知机，将解码器输出映射到锚框
- query_embed：创建查询嵌入，查询数量为num_queries、查询维度是hidden_dim
- input_proj：创建卷积层，将骨干网络的输出num_channels映射到hidden_dim，方便调用transformer
- backbone：骨干网络
- aux_loss：辅助解码损失函数
前向传播
- 首先利用isinstance函数判断并将输入的samples类型转换到NestedTensor
- 通过backbone网络提取特征，将特征图和位置编码分别存储到features和pos
- 将骨干网络提取到的features分解为原输入（针对transformer）和掩码
- 经过注意力模型获取输出
- 将多头注意力的输出通过嵌入层映射到类别
- 将多头注意力的输出通过卷积层归一化后映射到锚框预测
- 保存解码器输出层的预测概率和锚框。
- 根据aux_loss参数选择是否保存解码器中间层的预测概率和损失

segmentation

MHAttentionMap

class MHAttentionMap(nn.Module):
    """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)"""
    # 只返回注意力分数的softmax结果（即所有键的权重）
    def __init__(self, query_dim, hidden_dim, num_heads, dropout=0.0, bias=True):
        super().__init__()
        self.num_heads = num_heads
        self.hidden_dim = hidden_dim
        self.dropout = nn.Dropout(dropout)

        self.q_linear = nn.Linear(query_dim, hidden_dim, bias=bias)  # 查询的线性映射 将[n, query_dim] -> [n, hidden_dim]
        self.k_linear = nn.Linear(query_dim, hidden_dim, bias=bias)  # 键值的线性映射 将[n, query_dim] -> [n, hidden_dim]

        nn.init.zeros_(self.k_linear.bias)
        nn.init.zeros_(self.q_linear.bias)
        nn.init.xavier_uniform_(self.k_linear.weight)
        nn.init.xavier_uniform_(self.q_linear.weight)
        self.normalize_fact = float(hidden_dim / self.num_heads) ** -0.5  # 注意力缩放因子，用于缩放注意力分数，
                                                                          # 以确保不同头之间的计算稳定性。

    def forward(self, 
                q,  # [batchsize * seq_length, query_dim] 
                k,  # [batchsize, channels, h, w]
                mask: Optional[Tensor] = None):
        # 计算查询和键值之间的注意力权重
        q = self.q_linear(q)
        k = F.conv2d(k, self.k_linear.weight.unsqueeze(-1).unsqueeze(-1), self.k_linear.bias)  # 卷积核:[num_query, hidden_dim, 1, 1]
                                                                                               # 通过1x1卷积核实行线性变换
        qh = q.view(q.shape[0], q.shape[1], self.num_heads, self.hidden_dim // self.num_heads)  # [batch_size, seq_length, num_heads, head_dim]
        kh = k.view(k.shape[0], self.num_heads, self.hidden_dim // self.num_heads, k.shape[-2], k.shape[-1])  # [batch_size, num_heads, head_dim, h, w]
        # 对于每个batch_size和每个num_heads
        # qh 中的每个[seq_length, head_dim]与kh中的每个[head_dim, height, width]相乘
        # 并沿着head_dim求和，得到的结果形状为[batch_size, seq_length, num_heads, height, width]
        weights = torch.einsum("bqnc,bnchw->bqnhw",
                               qh * self.normalize_fact, # [batch_size, seq_length, num_heads, head_dim] 乘以 注意力缩放因子
                               kh)  # [batch_size, num_heads, head_dim, height, width]
        
        if mask is not None:
            weights.masked_fill_(mask.unsqueeze(1).unsqueeze(1), float("-inf"))  # 将掩码位置的权重设置为负无穷大，以避免这些位置的注意力。
        weights = F.softmax(weights.flatten(2),  #[batchsize, num_heads, head_dim, h * w]
                             dim=-1  # 在h * w维度上应用softmax
                             ).view(weights.size())  # 恢复形状为[batchsize, num_heads, head_dim, h, w]
        weights = self.dropout(weights)
        return weights

该类实现了一个注意力汇聚机制:

初始化:
- q_linear：查询的线性映射将[n, query_dim] -> [n, hidden_dim]
- k_linear键值的线性映射将[n, query_dim] -> [n, hidden_dim]
- 偏置初始化为0
- 权重矩阵使用xavier方法初始化
前向传播
- 输入中q的维度为[batchsize, num_queries, hidden_dim]，k的维度为[batchsize, hidden_dim, h, w]
- 首先利用线性层和和卷积层将q的维度变换为[batchsize, num_queries, hidden_dim] ，将k的维度变换为[batchsize, hidden_dim, h, w]
- 然后利用view函数进行维度变换，得到注意力输出qh形状为[batchsize, num_queries, num_heads, head_dim]，kh形状为[batch_size, num_heads, head_dim, h, w]
- 利用einsum得到最终点积结果，形状为[batch_size, seq_length, num_heads, height, width]，对应未经标准化的权重大小
- 当提供掩码时将掩码位置注意力权重设置为负无穷
- 经过softmax函数归一化得到加和为1的注意力权重,并通过dropout层返回

MaskHeadSmallConv

class MaskHeadSmallConv(nn.Module):
    """
    Simple convolutional head, using group norm.
    Upsampling is done using a FPN approach
    """

    def __init__(self, 
                 dim,  # 输入特征图的通道数。
                 fpn_dims,  # 表示来自特征金字塔网络的不同尺度的特征图的通道数
                 context_dim  # 上下文维度，用于定义中间层的通道数
                 ):
        super().__init__()
        # 包含了各个中间层的通道数，逐层减少 填充为1保持输出特征大小不变
        inter_dims = [dim, context_dim // 2, context_dim // 4, context_dim // 8, context_dim // 16, context_dim // 64]
        self.lay1 = torch.nn.Conv2d(dim, dim, 3, padding=1)
        self.gn1 = torch.nn.GroupNorm(8, dim)
        self.lay2 = torch.nn.Conv2d(dim, inter_dims[1], 3, padding=1)
        self.gn2 = torch.nn.GroupNorm(8, inter_dims[1])
        self.lay3 = torch.nn.Conv2d(inter_dims[1], inter_dims[2], 3, padding=1)
        self.gn3 = torch.nn.GroupNorm(8, inter_dims[2])
        self.lay4 = torch.nn.Conv2d(inter_dims[2], inter_dims[3], 3, padding=1)
        self.gn4 = torch.nn.GroupNorm(8, inter_dims[3])
        self.lay5 = torch.nn.Conv2d(inter_dims[3], inter_dims[4], 3, padding=1)
        self.gn5 = torch.nn.GroupNorm(8, inter_dims[4])
        self.out_lay = torch.nn.Conv2d(inter_dims[4], 1, 3, padding=1)

        self.dim = dim
        # 1x1卷积核,将通道数从fpn_dims[i]映射到inter_dims[i+1]
        self.adapter1 = torch.nn.Conv2d(fpn_dims[0],  # 1024 
                                        inter_dims[1], 1)
        self.adapter2 = torch.nn.Conv2d(fpn_dims[1],  # 512
                                        inter_dims[2], 1)
        self.adapter3 = torch.nn.Conv2d(fpn_dims[2],  # 256
                                        inter_dims[3], 1)

        for m in self.modules():  # 对所有卷积层初始化
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_uniform_(m.weight, a=1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: Tensor,  # [batchsize, channels, h, w]
                bbox_mask: Tensor,  # [batchsize, num_heads, head_dim, h, w]
                fpns: List[Tensor]
                ):
        x = torch.cat([_expand(x, bbox_mask.shape[1]), # [batchsize * num_heads, channels, h, w]
                       bbox_mask.flatten(0, 1)],  # [batchsize * num_heads, head_dim, h, w]
                       1)  # [batchsize * num_heads + channels(hidden_dim), head_dim, h, w]

        x = self.lay1(x)
        x = self.gn1(x)
        x = F.relu(x)
        x = self.lay2(x)
        x = self.gn2(x)
        x = F.relu(x)

        cur_fpn = self.adapter1(fpns[0])
        if cur_fpn.size(0) != x.size(0):  # 检查两个张量的batchsize是否相同
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        # 对x进行上采样,使他的空间尺寸([h, w])与cur_fpn一致
        # 实现了两个特征图的融合
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay3(x)
        x = self.gn3(x)
        x = F.relu(x)

        cur_fpn = self.adapter2(fpns[1])
        if cur_fpn.size(0) != x.size(0):
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay4(x)
        x = self.gn4(x)
        x = F.relu(x)

        cur_fpn = self.adapter3(fpns[2])
        if cur_fpn.size(0) != x.size(0):
            cur_fpn = _expand(cur_fpn, x.size(0) // cur_fpn.size(0))
        x = cur_fpn + F.interpolate(x, size=cur_fpn.shape[-2:], mode="nearest")
        x = self.lay5(x)
        x = self.gn5(x)
        x = F.relu(x)

        x = self.out_lay(x)
        return x

该类定义了一个特征金字塔网络，通过把不同大小的特征图加在一起来尽可能多的捕获不同大小的特征图，从而解决了特征捕获差的问题。

DETRsegm

class DETRsegm(nn.Module):
    def __init__(self, detr, freeze_detr=False):
        super().__init__()
        self.detr = detr

        if freeze_detr:  # 根据freeze_detr参数判断是否关闭反向传播（冻结模型）
            for p in self.parameters():
                p.requires_grad_(False)

        hidden_dim, nheads = detr.transformer.d_model, detr.transformer.nhead
        self.bbox_attention = MHAttentionMap(hidden_dim, hidden_dim, nheads, dropout=0.0)  # 用于处理边界框的注意力映射
        self.mask_head = MaskHeadSmallConv(hidden_dim + nheads, [1024, 512, 256], hidden_dim)  # 用于在分割任务中生成掩码

    def forward(self, samples: NestedTensor):
        if isinstance(samples, (list, torch.Tensor)):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.detr.backbone(samples)  # 输入通过网络骨架得到特征图和位置编码

        bs = features[-1].tensors.shape[0]  # 获取transformer输入的批量大小

        src, mask = features[-1].decompose()  # 将特征图拆分为transformer的输入和掩码
        assert mask is not None
        src_proj = self.detr.input_proj(src)  # 将输入通道转换为隐藏层维度,以适应transformer输入格式
        # hs:解码器的输出  memory:编码器的输出
        hs, memory = self.detr.transformer(src_proj, mask, self.detr.query_embed.weight, pos[-1])

        outputs_class = self.detr.class_embed(hs)  # 将transformer的输出映射到类别上
        outputs_coord = self.detr.bbox_embed(hs).sigmoid()  # 将transformer的输出映射到多层感知机上,用于预测锚框位置
        out = {"pred_logits": outputs_class[-1], "pred_boxes": outputs_coord[-1]}
        if self.detr.aux_loss:
            out['aux_outputs'] = self.detr._set_aux_loss(outputs_class, outputs_coord)

        # FIXME h_boxes takes the last one computed, keep this in mind
        bbox_mask = self.bbox_attention(hs[-1], memory, mask=mask)  # 用于确定模型应该关注图像中的哪些区域来预测边界框。
                                                                    # [batchsize, num_heads, head_dim, h, w]
        seg_masks = self.mask_head(src_proj, bbox_mask, [features[2].tensors,  # 传参传的是out,即特征图
                                                          features[1].tensors, features[0].tensors])  # 返回一个可以捕获不同尺寸的特征的特征金字塔
        outputs_seg_masks = seg_masks.view(bs, self.detr.num_queries, seg_masks.shape[-2], seg_masks.shape[-1])  # 将特征图形状调整为[batch_size, num_queries, h, w]

        out["pred_masks"] = outputs_seg_masks
        return out

该类通过bbox_attention模块，利用注意力机制来帮助模型在解码器输出和编码器记忆之间进行信息交互，并利用mask_head模块生成分割的掩码，用来确定模型应该关注哪一些特征来预测边界框。

y0k1n0

关注

11
点赞
踩
23

收藏

觉得还不错? 一键收藏
0
评论
detr源码解读

这段代码利用掩码的累加作为位置编码的posPEpos2isin10000dmodel2iposPEpos2i1cos10000dmodel2ipos其中，pos表示token在序列中的位置，设句子长度为 L，则pos01L−1是token的位置向量，PEpos2i表示这个位置向量里的第i个元素，2i1表示奇数维度，2i表示偶数维度；dmodel。
复制链接

扫一扫

专栏目录