Transformer实战-系列教程16：DETR 源码解读3（DETR类）

Transformer实战-系列教程总目录

有任何问题欢迎在下面留言
本篇文章的代码运行界面均在Pycharm中进行
本篇文章配套的代码资源已经上传
点我下载源码

DETR 算法解读
DETR 源码解读1（项目配置/CocoDetection类）
DETR 源码解读2（ConvertCocoPolysToMask类）
DETR 源码解读3（DETR类）
DETR 源码解读4（Joiner类/PositionEmbeddingSine类/位置编码/backbone）

4、DETR类

位置：models/detr.py/DETR类

4.1 构造函数

class DETR(nn.Module):
    def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
        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 = 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

1. DETR类继承torch nn.Module
2. 构造函数，传入5个参数：
   • backbone：CNN骨架网络，用于特征提取
   • transformer：Transformer模型，用于处理序列数据
   • num_classes：目标类别的数量
   • num_queries：解码器初始化生成的100个向量的个数，num_queries=100
   • aux_loss：一个布尔值，指示是否使用辅助损失来帮助训练
3. 初始化
4. num_queries
5. transformer
6. hidden_dim ，Transformer中的隐藏层维度
7. class_embed ，类别预测的输出层，这个全连接层是接Transformer的输出，类别加1是额外的无类别对象
8. bbox_embed，一个MLP，也是接Transformer的输出，边界框的四个坐标的回归
9. query_embed ，解码器的初始100个向量
10. input_proj ，一个1x1的二维卷积，使得backbone的输出通道数映射到与Transformer隐藏层维度相同
11. backbone，一个预训练的卷积神经网络，主要作用是提取图像的特征，它的输出经过input_proj 处理后作为Transformer的输入
12. aux_loss，保存是否使用辅助损失的标志

这里包含了几个自定义函数和类：
nested_tensor_from_tensor_list函数：将不同尺寸处理的图像Tensor转换为一个嵌套Tensor
MLP类：边界框的四个坐标的回归
transformer类：构建transformer架构
backbone：用于提取图像特征的CNN

4.2 前向传播

    def forward(self, samples: NestedTensor):
        if isinstance(samples, (list, torch.Tensor)):
            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), mask, self.query_embed.weight, pos[-1])[0]
        outputs_class = self.class_embed(hs)
        outputs_coord = self.bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        if self.aux_loss:
            out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
        return out    

1. 前向传播函数，输入为samples=NestedTensor{mask={Tensor(2,771,911)}，tensors={Tensor(2,3,771,911)}}
2. 检查samples是否为列表或Tensor类型
3. samples ，如果是，使用nested_tensor_from_tensor_list函数转换为NestedTensor
4. features, pos，图像特征图列表和对应的位置编码列表，backbone实际上一个现在的resnet
5. src, mask，解构最后一层的特征，获取源数据和掩码，src：torch.Size([2, 2048, 21, 18])，mask torch.Size([2, 21, 18])，2是batch，2048是特征维度，后面两个是图像长宽
6. 确保掩码不为空
7. 将数据通过Transformer处理，获取序列输出，torch.Size([6, 2, 100, 256])，6是Transformer的堆叠层数，2是batch，100是生成100个目标预测，256是每个目标预测的维度
8. outputs_class ，获取类别预测
9. outputs_coord ，获取边界框坐标预测，并使用sigmoid函数将输出值限制在0到1之间
10. out ，将类别预测结果和 边界框坐标预测结果做成一个字典
11. 如果启用了辅助损失
12. 通过辅助函数_set_aux_loss计算辅助损失
13. 返回out

4.3 辅助函数_set_aux_loss()

@torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

1. @torch.jit.unused：一个装饰器，指示当使用TorchScript编译模型时，该方法不应被编译。这是因为辅助损失的计算可能不兼容TorchScript的静态图特性
2. 定义函数，接收类别预测和边界框坐标作为输入
3. 返回一个列表，将每一个类别预测和边界框坐标都封装成一个字典，这样，训练过程中可以计算每一层的损失，从而实现辅助损失的目的

DETR 算法解读
DETR 源码解读1（项目配置/CocoDetection类）
DETR 源码解读2（ConvertCocoPolysToMask类）
DETR 源码解读3（DETR类）
DETR 源码解读4（Joiner类/PositionEmbeddingSine类/位置编码/backbone）