DETR的一个简单实现Demo

这个demo会基于预训练权重实现一个DETR的简化版,然后对一张图片作预测,最后展示出预测效果。

首先导入需要的相关库:

from PIL import Image  
import requests  
import matplotlib.pyplot as plt  
%config InlineBackend.figure_format = 'retina'  
  
import torch  
from torch import nn  
from torchvision.models import resnet50  
import torchvision.transforms as T  
  
# 推理过程,不需要梯度  
torch.set_grad_enabled(False);

先展示demo模型的完整代码

class DETRdemo(nn.Module):
    """
    DETR的一个基本实现

    此演示Demo与论文中的完整DETR模型有一下不同点:
    * 使用的是可学习位置编码(论文中使用的是正余弦位置编码)
    * 位置编码在输入端传入 (原文在注意力模块传入)
    * 采用fc bbox 预测器 (没有采用MLP)
    该模型在 COCO val5k 上达到约 40 AP,在 Tesla V100 上以约 28 FPS 的速度运行。
    仅支持批量大小为1。
    """
    def __init__(self, num_classes, hidden_dim=256, nheads=8,
                 num_encoder_layers=6, num_decoder_layers=6):
        # hidden_dim: 隐藏状态的神经单元个数,也就是隐藏层的节点数,应该可以按计算需要“随意”设置。
        super().__init__()

        # create ResNet-50 backbone
        # 创建Resnet50
        self.backbone = resnet50()
        # 删除最后的全连接层
        del self.backbone.fc

        # create conversion layer
        # 将骨干网络的输出特征图维度映射到Transformer输入所需的维度
        self.conv = nn.Conv2d(2048, hidden_dim, 1)

        # create a default PyTorch transformer
        # nheads代表多头注意力的"头数"
        self.transformer = nn.Transformer(
            hidden_dim, nheads, num_encoder_layers, num_decoder_layers)

        # prediction heads, one extra class for predicting non-empty slots
        # num_classes需要在原有来别数量上多加一个non-empty类
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        # note that in baseline DETR linear_bbox layer is 3-layer MLP
        # 标准DETR模型中最后的输出层由三个全连接层构成而非一个全连接层
        # bbox的形式是(x,y,w,h),因此是四维
        self.linear_bbox = nn.Linear(hidden_dim, 4)

        # output positional encodings (object queries)
        # 用于解码器输入的位置编码,100代表最终解码出100个物体
        # 即对一张图片(最多)检测出100个物体
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))

        # spatial positional encodings
        # note that in baseline DETR we use sine positional encodings
        # 用于编码器输入的位置编码
        # 对特征图的行、列分别进行位置编码,而后会将两者结果拼接
        # 因此维度格式hidden_dim的一半,前128是x后128是y
        # nn.Parameter() 在指定Tensor中随机生成参数
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))

    def forward(self, inputs):
        # propagate inputs through ResNet-50 up to avg-pool layer
        x = self.backbone.conv1(inputs)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)

        x = self.backbone.layer1(x)
        x = self.backbone.layer2(x)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)

        # convert from 2048 to 256 feature planes for the transformer
        # 将backbone的输出维度转换为Transformer输入所需的维度
        # h.shape = (1, hidden_dim, H, W)
        h = self.conv(x)


        # construct positional encodings
        H, W = h.shape[-2:]
        # Tensor.unsqueeze() 在指定位置插入新维度
        # Tensor.repeat() 沿某个维度复制Tensor
        # self.col_embed[:W].shape = (W, hidden_dim / 2) hidden_dim &#
  • 1
    点赞
  • 5
    收藏
    觉得还不错? 一键收藏
  • 3
    评论
Transformer在许多NLP(自然语言处理)任务中取得了最先进的成果。 DETR(Detection Transformer)是Facebook提出的基于Transformer的端到端目标检测方法。DETR使用CNN+Transformer进行图像目标检测,该方法没有NMS后处理步骤、没有anchor。DETR总体思路是把目标检测看成一个set prediction的问题,并且使用Transformer来预测物体边界框的集合。本课程对DETR的原理与PyTorch实现代码进行精讲,来帮助大家掌握其详细原理和具体实现。 原理精讲部分包括:Transformer的架构概述、Transformer的Encoder 、Transformer的Decoder、DETR网络架构、DETR损失函数、DETR实验结果和分析。  代码精讲部分使用Jupyter Notebook对DETR的PyTorch代码进行逐行解读,包括:安装PyTorch、 DETR官方DemoDETR的hands-on tutorial,DETR的代码精讲(数据准备、Backbone和位置编码、Transformer架构的实现)。相关课程: 《Transformer原理与代码精讲(PyTorch)》https://edu.csdn.net/course/detail/36697《Transformer原理与代码精讲(TensorFlow)》https://edu.csdn.net/course/detail/36699《ViT(Vision Transformer)原理与代码精讲》https://edu.csdn.net/course/detail/36719《DETR原理与代码精讲》https://edu.csdn.net/course/detail/36768《Swin Transformer实战目标检测:训练自己的数据集》https://edu.csdn.net/course/detail/36585《Swin Transformer实战实例分割:训练自己的数据集》https://edu.csdn.net/course/detail/36586《Swin Transformer原理与代码精讲》 https://download.csdn.net/course/detail/37045
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值