DETR的一个简单实现Demo

该博客通过一个DETR的简化版demo,演示如何基于预训练权重进行目标检测,包括模型结构、位置编码、Transformer的使用,以及预测结果的处理和可视化。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

这个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() 在指定位置插入新维度
        
实现基于 DETR 的实例分割模型时,通常需要对原始 DETR 模型进行扩展,使其能够输出像素级的掩码。DETR 本身主要关注目标检测任务中的边界框和类别预测,但其结构为引入额外的解码头提供了良好的基础。 ### 数据准备 首先,确保使用支持实例分割的数据集,如 COCO。数据加载器需要同时提供图像和对应的实例掩码标签。可以参考以下代码片段构建数据集: ```python from torch.utils.data import DataLoader from torchvision import transforms from torchvision.datasets import CocoDetection transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = CocoDetection(root='path/to/train/images', annFile='path/to/annotations/instances_train.json', transform=transform) train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=4) ``` ### 构建 DETR 模型 PyTorch 提供了 DETR 的预训练模型接口,可以通过 `torchvision.models.detection` 加载,并在此基础上添加掩码预测模块。以下是一个简化版的 DETR 实例分割模型构建示例: ```python import torch import torchvision from torchvision.models.detection.detr import DETR # 定义自定义的DETR模型,增加掩码分支 class DETRInstanceSegmentation(DETR): def __init__(self, num_classes, hidden_dim=256, nheads=8, num_encoder_layers=6, num_decoder_layers=6): super().__init__( backbone=ResNetBackbone(), # 自定义的backbone transformer=Transformer(nhead=nheads, num_encoder_layers=num_encoder_layers, num_decoder_layers=num_decoder_layers), num_classes=num_classes, hidden_dim=hidden_dim ) self.mask_head = torch.nn.Sequential( torch.nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1), torch.nn.ReLU(), torch.nn.Conv2d(hidden_dim, 1, kernel_size=1) # 输出单通道掩码 ) def forward(self, images): features = self.backbone(images) out = self.transformer(features) mask_logits = self.mask_head(out.unsqueeze(-1).unsqueeze(-1)) # 扩展维度以适配卷积层 return out, mask_logits.squeeze() ``` ### 训练模型 在训练过程中,除了原有的分类损失和边界框回归损失外,还需要加入掩码预测的损失函数,例如二元交叉熵损失(BCE Loss): ```python model = DETRInstanceSegmentation(num_classes=91) # 包括背景类 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) criterion_class = torch.nn.CrossEntropyLoss() criterion_bbox = torch.nn.L1Loss() criterion_mask = torch.nn.BCEWithLogitsLoss() for epoch in range(10): # 进行多轮训练 model.train() for images, targets in train_loader: images = list(image.to(device) for image in images) labels = [t['labels'].to(device) for t in targets] boxes = [t['boxes'].to(device) for t in targets] masks = [t['masks'].float().to(device) for t in targets] # 转换为浮点数 class_outputs, bbox_outputs, mask_outputs = model(images) loss_class = criterion_class(class_outputs, torch.cat(labels)) loss_bbox = criterion_bbox(bbox_outputs, torch.cat(boxes)) loss_mask = criterion_mask(mask_outputs, torch.cat(masks)) total_loss = loss_class + loss_bbox + loss_mask optimizer.zero_grad() total_loss.backward() optimizer.step() print(f"Epoch {epoch} Loss: {total_loss.item()}") ``` ### 测试与推理 测试阶段,模型会生成边界框、类别和掩码预测结果。可以使用如下代码进行可视化: ```python model.eval() with torch.no_grad(): for images, _ in test_loader: images = list(image.to(device) for image in images) class_preds, bbox_preds, mask_preds = model(images) # 可视化第一个样本的结果 image = images[0].cpu().permute(1, 2, 0).numpy() mask_pred = torch.sigmoid(mask_preds[0]).cpu().numpy() plt.figure(figsize=(10, 5)) plt.subplot(1, 2, 1) plt.imshow(image) plt.title("Input Image") plt.subplot(1, 2, 2) plt.imshow(mask_pred, cmap='viridis') plt.title("Predicted Mask") plt.show() ``` ###
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值