devbox
从前慢现在也慢
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【YOLOv8改进[Head检测头]】YOLOv8换个RT-DETR head助力模型更优秀

作者：从前慢现在也慢 | 2024-04-24 10:27:33

赞

踩

【YOLOv8改进[Head检测头]】YOLOv8换个RT-DETR head助力模型更优秀

一RT-DETR

官方论文地址：https://arxiv.org/pdf/2304.08069.pdf

因为YOLO的合理速度和准确性之间的权衡, 这一系列已成为最流行的实时目标检测框架。然而，观察到nms对yolo的速度和准确性产生了负面影响。最近，基于端到端变换器的检测器(DETRs)消除了传统实时检测器中的非最大抑制（NMS）等后处理步骤的需要，这些步骤一直是传统实时检测器中的瓶颈，提供了一种替代方案。然而，高昂的计算成本限制了它们的实用性，阻碍了它们充分发挥不用NMS的优势。

在本文中，提出了实时检测变换器(RT-DETR)，这是我们所知的第一个解决上述困境的实时端到端目标检测器。在先进的DETR基础上分两步构建RT-DETR:首先专注于在提高速度的同时保持精度，其次是在提高精度的同时保持速度。具体而言，设计了一种高效的混合编码器，通过解耦尺度内相互作用和跨尺度融合来快速处理多尺度特征，以提高速度。然后，提出了不确定性最小的查询选择，为解码器提供高质量的初始查询，从而提高准确率。此外，RT-DETR支持灵活的速度调整，通过调整解码器层的数量，以适应各种场景，而无需重新训练。

RT-DETR-R50/ R101在COCO上实现53.1% / 54.3%的AP，在T4 GPU上实现108 / 74 FPS，在速度和精度方面都优于以前先进的yolo。此外，RT-DETR-R50在精度上比DINO-R50高出2.2%，在FPS上高出约21倍。RT - DETR - R50 / R101经过Objects365预训练，AP达到55.3% / 56.2%。

官方代码地址：DETRs Beat YOLOs on Real-time Object Detection

综上，RT-DETR模型建立在于两个关键创新：

① 高效混合编码器：通过解耦内部尺度交互和跨尺度融合来处理多尺度特征。这种设计显著降低了计算负担，同时保持了高性能，实现了实时目标检测。

② 提出了不确定性最小的查询选择，为解码器提供高质量的初始查询，从而提高准确率。

1 编码器结构

下图是每个变体的编码器结构。SSE表示单尺度Transformer编码器，MSE表示多尺度Transformer编码器，CSF表示跨尺度融合。AIFI和CCFF是我们设计的混合编码器的两个模块。

2 RT-DETR

下图为RT-DETR概述。将主干最后三个阶段的特征输入到编码器中。高效混合编码器通过基于注意力的尺度内特征交互(AIFI)和基于cnn的跨尺度特征融合(CCFF)将多尺度特征转化为图像特征序列。然后，最小不确定性查询选择固定数量的编码器特征作为解码器的初始对象查询。最后，具有辅助预测头的解码器迭代优化对象查询以生成类别和框。

3 CCFF中的融合块

下图为 CCFF中的融合块。

实验结果：

二 RT-DETR检测头的代码

RT-DETR检测头的代码如下所示：

  1. class RTDETRDecoder(nn.Module):
  2.     """
  3.     Real-Time Deformable Transformer Decoder (RTDETRDecoder) module for object detection.
  5.     This decoder module utilizes Transformer architecture along with deformable convolutions to predict bounding boxes
  6.     and class labels for objects in an image. It integrates features from multiple layers and runs through a series of
  7.     Transformer decoder layers to output the final predictions.
  8.     """
  9.  
  10.     export = False  # export mode
  11.  
  12.     def __init__(
  13.         self,
  14.         nc=80,
  15.         ch=(512, 1024, 2048),
  16.         hd=256,  # hidden dim
  17.         nq=300,  # num queries
  18.         ndp=4,  # num decoder points
  19.         nh=8,  # num head
  20.         ndl=6,  # num decoder layers
  21.         d_ffn=1024,  # dim of feedforward
  22.         dropout=0.0,
  23.         act=nn.ReLU(),
  24.         eval_idx=-1,
  25.         # Training args
  26.         nd=100,  # num denoising
  27.         label_noise_ratio=0.5,
  28.         box_noise_scale=1.0,
  29.         learnt_init_query=False,
  30.     ):
  31.         """
  32.         Initializes the RTDETRDecoder module with the given parameters.
  34.         Args:
  35.             nc (int): Number of classes. Default is 80.
  36.             ch (tuple): Channels in the backbone feature maps. Default is (512, 1024, 2048).
  37.             hd (int): Dimension of hidden layers. Default is 256.
  38.             nq (int): Number of query points. Default is 300.
  39.             ndp (int): Number of decoder points. Default is 4.
  40.             nh (int): Number of heads in multi-head attention. Default is 8.
  41.             ndl (int): Number of decoder layers. Default is 6.
  42.             d_ffn (int): Dimension of the feed-forward networks. Default is 1024.
  43.             dropout (float): Dropout rate. Default is 0.
  44.             act (nn.Module): Activation function. Default is nn.ReLU.
  45.             eval_idx (int): Evaluation index. Default is -1.
  46.             nd (int): Number of denoising. Default is 100.
  47.             label_noise_ratio (float): Label noise ratio. Default is 0.5.
  48.             box_noise_scale (float): Box noise scale. Default is 1.0.
  49.             learnt_init_query (bool): Whether to learn initial query embeddings. Default is False.
  50.         """
  51.         super().__init__()
  52.         self.hidden_dim = hd
  53.         self.nhead = nh
  54.         self.nl = len(ch)  # num level
  55.         self.nc = nc
  56.         self.num_queries = nq
  57.         self.num_decoder_layers = ndl
  58.  
  59.         # Backbone feature projection
  60.         self.input_proj = nn.ModuleList(nn.Sequential(nn.Conv2d(x, hd, 1, bias=False), nn.BatchNorm2d(hd)) for x in ch)
  61.         # NOTE: simplified version but it's not consistent with .pt weights.
  62.         # self.input_proj = nn.ModuleList(Conv(x, hd, act=False) for x in ch)
  63.  
  64.         # Transformer module
  65.         decoder_layer = DeformableTransformerDecoderLayer(hd, nh, d_ffn, dropout, act, self.nl, ndp)
  66.         self.decoder = DeformableTransformerDecoder(hd, decoder_layer, ndl, eval_idx)
  67.  
  68.         # Denoising part
  69.         self.denoising_class_embed = nn.Embedding(nc, hd)
  70.         self.num_denoising = nd
  71.         self.label_noise_ratio = label_noise_ratio
  72.         self.box_noise_scale = box_noise_scale
  73.  
  74.         # Decoder embedding
  75.         self.learnt_init_query = learnt_init_query
  76.         if learnt_init_query:
  77.             self.tgt_embed = nn.Embedding(nq, hd)
  78.         self.query_pos_head = MLP(4, 2 * hd, hd, num_layers=2)
  79.  
  80.         # Encoder head
  81.         self.enc_output = nn.Sequential(nn.Linear(hd, hd), nn.LayerNorm(hd))
  82.         self.enc_score_head = nn.Linear(hd, nc)
  83.         self.enc_bbox_head = MLP(hd, hd, 4, num_layers=3)
  84.  
  85.         # Decoder head
  86.         self.dec_score_head = nn.ModuleList([nn.Linear(hd, nc) for _ in range(ndl)])
  87.         self.dec_bbox_head = nn.ModuleList([MLP(hd, hd, 4, num_layers=3) for _ in range(ndl)])
  88.  
  89.         self._reset_parameters()
  90.  
  91.     def forward(self, x, batch=None):
  92.         """Runs the forward pass of the module, returning bounding box and classification scores for the input."""
  93.         from ultralytics.models.utils.ops import get_cdn_group
  94.  
  95.         # Input projection and embedding
  96.         feats, shapes = self._get_encoder_input(x)
  97.  
  98.         # Prepare denoising training
  99.         dn_embed, dn_bbox, attn_mask, dn_meta = get_cdn_group(
  100.             batch,
  101.             self.nc,
  102.             self.num_queries,
  103.             self.denoising_class_embed.weight,
  104.             self.num_denoising,
  105.             self.label_noise_ratio,
  106.             self.box_noise_scale,
  107.             self.training,
  108.         )
  109.  
  110.         embed, refer_bbox, enc_bboxes, enc_scores = self._get_decoder_input(feats, shapes, dn_embed, dn_bbox)
  111.  
  112.         # Decoder
  113.         dec_bboxes, dec_scores = self.decoder(
  114.             embed,
  115.             refer_bbox,
  116.             feats,
  117.             shapes,
  118.             self.dec_bbox_head,
  119.             self.dec_score_head,
  120.             self.query_pos_head,
  121.             attn_mask=attn_mask,
  122.         )
  123.         x = dec_bboxes, dec_scores, enc_bboxes, enc_scores, dn_meta
  124.         if self.training:
  125.             return x
  126.         # (bs, 300, 4+nc)
  127.         y = torch.cat((dec_bboxes.squeeze(0), dec_scores.squeeze(0).sigmoid()), -1)
  128.         return y if self.export else (y, x)
  129.  
  130.     def _generate_anchors(self, shapes, grid_size=0.05, dtype=torch.float32, device="cpu", eps=1e-2):
  131.         """Generates anchor bounding boxes for given shapes with specific grid size and validates them."""
  132.         anchors = []
  133.         for i, (h, w) in enumerate(shapes):
  134.             sy = torch.arange(end=h, dtype=dtype, device=device)
  135.             sx = torch.arange(end=w, dtype=dtype, device=device)
  136.             grid_y, grid_x = torch.meshgrid(sy, sx, indexing="ij") if TORCH_1_10 else torch.meshgrid(sy, sx)
  137.             grid_xy = torch.stack([grid_x, grid_y], -1)  # (h, w, 2)
  138.  
  139.             valid_WH = torch.tensor([w, h], dtype=dtype, device=device)
  140.             grid_xy = (grid_xy.unsqueeze(0) + 0.5) / valid_WH  # (1, h, w, 2)
  141.             wh = torch.ones_like(grid_xy, dtype=dtype, device=device) * grid_size * (2.0**i)
  142.             anchors.append(torch.cat([grid_xy, wh], -1).view(-1, h * w, 4))  # (1, h*w, 4)
  143.  
  144.         anchors = torch.cat(anchors, 1)  # (1, h*w*nl, 4)
  145.         valid_mask = ((anchors > eps) & (anchors < 1 - eps)).all(-1, keepdim=True)  # 1, h*w*nl, 1
  146.         anchors = torch.log(anchors / (1 - anchors))
  147.         anchors = anchors.masked_fill(~valid_mask, float("inf"))
  148.         return anchors, valid_mask
  149.  
  150.     def _get_encoder_input(self, x):
  151.         """Processes and returns encoder inputs by getting projection features from input and concatenating them."""
  152.         # Get projection features
  153.         x = [self.input_proj[i](feat) for i, feat in enumerate(x)]
  154.         # Get encoder inputs
  155.         feats = []
  156.         shapes = []
  157.         for feat in x:
  158.             h, w = feat.shape[2:]
  159.             # [b, c, h, w] -> [b, h*w, c]
  160.             feats.append(feat.flatten(2).permute(0, 2, 1))
  161.             # [nl, 2]
  162.             shapes.append([h, w])
  163.  
  164.         # [b, h*w, c]
  165.         feats = torch.cat(feats, 1)
  166.         return feats, shapes
  167.  
  168.     def _get_decoder_input(self, feats, shapes, dn_embed=None, dn_bbox=None):
  169.         """Generates and prepares the input required for the decoder from the provided features and shapes."""
  170.         bs = feats.shape[0]
  171.         # Prepare input for decoder
  172.         anchors, valid_mask = self._generate_anchors(shapes, dtype=feats.dtype, device=feats.device)
  173.         features = self.enc_output(valid_mask * feats)  # bs, h*w, 256
  174.  
  175.         enc_outputs_scores = self.enc_score_head(features)  # (bs, h*w, nc)
  176.  
  177.         # Query selection
  178.         # (bs, num_queries)
  179.         topk_ind = torch.topk(enc_outputs_scores.max(-1).values, self.num_queries, dim=1).indices.view(-1)
  180.         # (bs, num_queries)
  181.         batch_ind = torch.arange(end=bs, dtype=topk_ind.dtype).unsqueeze(-1).repeat(1, self.num_queries).view(-1)
  182.  
  183.         # (bs, num_queries, 256)
  184.         top_k_features = features[batch_ind, topk_ind].view(bs, self.num_queries, -1)
  185.         # (bs, num_queries, 4)
  186.         top_k_anchors = anchors[:, topk_ind].view(bs, self.num_queries, -1)
  187.  
  188.         # Dynamic anchors + static content
  189.         refer_bbox = self.enc_bbox_head(top_k_features) + top_k_anchors
  190.  
  191.         enc_bboxes = refer_bbox.sigmoid()
  192.         if dn_bbox is not None:
  193.             refer_bbox = torch.cat([dn_bbox, refer_bbox], 1)
  194.         enc_scores = enc_outputs_scores[batch_ind, topk_ind].view(bs, self.num_queries, -1)
  195.  
  196.         embeddings = self.tgt_embed.weight.unsqueeze(0).repeat(bs, 1, 1) if self.learnt_init_query else top_k_features
  197.         if self.training:
  198.             refer_bbox = refer_bbox.detach()
  199.             if not self.learnt_init_query:
  200.                 embeddings = embeddings.detach()
  201.         if dn_embed is not None:
  202.             embeddings = torch.cat([dn_embed, embeddings], 1)
  203.  
  204.         return embeddings, refer_bbox, enc_bboxes, enc_scores
  205.  
  206.     # TODO
  207.     def _reset_parameters(self):
  208.         """Initializes or resets the parameters of the model's various components with predefined weights and biases."""
  209.         # Class and bbox head init
  210.         bias_cls = bias_init_with_prob(0.01) / 80 * self.nc
  211.         # NOTE: the weight initialization in `linear_init` would cause NaN when training with custom datasets.
  212.         # linear_init(self.enc_score_head)
  213.         constant_(self.enc_score_head.bias, bias_cls)
  214.         constant_(self.enc_bbox_head.layers[-1].weight, 0.0)
  215.         constant_(self.enc_bbox_head.layers[-1].bias, 0.0)
  216.         for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head):
  217.             # linear_init(cls_)
  218.             constant_(cls_.bias, bias_cls)
  219.             constant_(reg_.layers[-1].weight, 0.0)
  220.             constant_(reg_.layers[-1].bias, 0.0)
  221.  
  222.         linear_init(self.enc_output[0])
  223.         xavier_uniform_(self.enc_output[0].weight)
  224.         if self.learnt_init_query:
  225.             xavier_uniform_(self.tgt_embed.weight)
  226.         xavier_uniform_(self.query_pos_head.layers[0].weight)
  227.         xavier_uniform_(self.query_pos_head.layers[1].weight)
  228.         for layer in self.input_proj:
  229.             xavier_uniform_(layer[0].weight)

三 YOLOv8换个RT-DETR head

ultralytics的版本为8.1.47，如下图所示：

1 总体修改

RT-DETR检测头已经集成在YOLOv8的项目里面了，我们可以直接使用。

注意：使用了RT-DETR检测头后，需要增加epoch。

2 配置文件

yolov8_RT-DETR.yaml的内容如下所示：

  1. # Ultralytics YOLO 
    声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/从前慢现在也慢/article/detail/478922
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。