Transformer实战-系列教程15:DETR 源码解读2(整体架构:DETR类)

🚩🚩🚩Transformer实战-系列教程总目录

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

DETR 算法解读
DETR 源码解读1(项目配置/CocoDetection类/ConvertCocoPolysToMask类)
DETR 源码解读2(DETR类)
DETR 源码解读3(位置编码:Joiner类/PositionEmbeddingSine类)
DETR 源码解读4(BackboneBase类/Backbone类)
DETR 源码解读5(Transformer类)
DETR 源码解读6(编码器:TransformerEncoder类/TransformerEncoderLayer类)
DETR 源码解读7(解码器:TransformerDecoder类/TransformerDecoderLayer类)
DETR 源码解读8(训练函数/损失函数)

4、DETR类

位置:models/detr.py/DETR类

4.1 构造函数

python 复制代码
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 前向传播

python 复制代码
    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,features和pos是一个list结构,保存了各层的输出
  5. src, mask,解构最后一层的特征,获取源数据和掩码,src:torch.Size([2, 2048, 21, 18]),mask torch.Size([2, 21, 18]),2是batch,2048是特征维度,后面两个是图像长宽,这里的features[-1]表示在backbone中有多层都有输出,features保存了各层的输出,这里-1就表示最后的输出
  6. 确保掩码不为空
  7. 将数据通过Transformer处理,获取序列输出,torch.Size([6, 2, 100, 256]),6是Transformer的堆叠层数,2是batch,100是生成100个目标预测,256是每个目标预测的维度,Transformer模块有两个返回值,只取第一个返回值
  8. outputs_class ,获取类别预测
  9. outputs_coord ,获取边界框坐标预测,并使用sigmoid函数将输出值限制在0到1之间
  10. out ,将类别预测结果和 边界框坐标预测结果做成一个字典
  11. 如果启用了辅助损失
  12. 通过辅助函数_set_aux_loss计算辅助损失
  13. 返回out

4.3 辅助函数_set_aux_loss()

python 复制代码
@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类/ConvertCocoPolysToMask类)
DETR 源码解读2(DETR类)
DETR 源码解读3(位置编码:Joiner类/PositionEmbeddingSine类)
DETR 源码解读4(BackboneBase类/Backbone类)
DETR 源码解读5(Transformer类)
DETR 源码解读6(编码器:TransformerEncoder类/TransformerEncoderLayer类)
DETR 源码解读7(解码器:TransformerDecoder类/TransformerDecoderLayer类)
DETR 源码解读8(训练函数/损失函数)

相关推荐
轻口味1 小时前
【每日学点鸿蒙知识】沙箱目录、图片压缩、characteristicsArray、gm-crypto 国密加解密、通知权限
pytorch·华为·harmonyos
Captain823Jack2 小时前
nlp新词发现——浅析 TF·IDF
人工智能·python·深度学习·神经网络·算法·自然语言处理
Captain823Jack3 小时前
w04_nlp大模型训练·中文分词
人工智能·python·深度学习·神经网络·算法·自然语言处理·中文分词
Wishell20153 小时前
Pytorch文件夹结构
pytorch
itwangyang5203 小时前
AIDD - 从机器学习到深度学习:蛋白质-配体对接评分函数的进展
人工智能·深度学习·机器学习
mingo_敏4 小时前
深度学习中的并行策略概述:2 Data Parallelism
人工智能·深度学习
湫ccc6 小时前
《Opencv》基础操作详解(3)
人工智能·opencv·计算机视觉
Jack_pirate6 小时前
深度学习中的特征到底是什么?
人工智能·深度学习
微凉的衣柜6 小时前
微软在AI时代的战略布局和挑战
人工智能·深度学习·microsoft
哦哦~9217 小时前
深度学习驱动的油气开发技术与应用
大数据·人工智能·深度学习·学习