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

【YOLOv8改进[检测头Head]】YOLOv8的“新头”之动态头（DynamicHead）

作者：weixin_40725706 | 2024-04-23 18:04:32

赞

踩

【YOLOv8改进[检测头Head]】YOLOv8的“新头”之动态头（DynamicHead）

目录

一 DynamicHead

二 YOLOv8的“新头”之动态头

1 总体修改

2 配置文件

3 训练

其他

一 DynamicHead

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

官方代码地址：GitCode - 开发者的代码家园

在计算机视觉应用中，目标检测是为了回答“是什么物体位于哪里”的问题的任务。如何提高目标检测头的性能已成为现有目标检测工作中的一个关键问题。

在本文中，提出了一种新的目标检测头，它将尺度感知、空间感知和任务感知的注意统一在一个框架中。提出了一种关注目标检测头的新视角。作为一个插件块，动态头可以灵活地集成到任何现有的目标检测器框架中，并在目标检测头中有效地应用了注意力机制，以提高其性能和效率。

下图说明了的动态头部方法。它包含三种不同的注意机制，每一种机制都侧重于不同的视角:尺度感知注意力、空间感知注意力和任务感知注意力。同时，还可视化了在每个注意模块之后特征图是如何改进的。

下图是动态头的详细设计。

(a) 显示了每个注意力模块的详细实施情况。

(b) 展示了如何将动态头部块应用于one-stage目标检测。

(c) 展示了如何将动态头部块应用于two-stage目标检测。

二 YOLOv8的“新头”之动态头

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

1 总体修改

① 添加DynamicHead.py文件

在ultralytics/nn/modules目录下新建DynamicHead.py文件，内容如下所示：

  1. import torch
  2. import math
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. from mmcv.ops import ModulatedDeformConv2d
  6. from ultralytics.utils.tal import dist2bbox, make_anchors
  7.  
  8. __all__ = ['Detect_DynamicHead']
  9.  
  10. def _make_divisible(v, divisor, min_value=None):
  11.     if min_value is None:
  12.         min_value = divisor
  13.     new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
  14.     # Make sure that round down does not go down by more than 10%.
  15.     if new_v < 0.9 * v:
  16.         new_v += divisor
  17.     return new_v
  18.  
  19.  
  20. class h_swish(nn.Module):
  21.     def __init__(self, inplace=False):
  22.         super(h_swish, self).__init__()
  23.         self.inplace = inplace
  24.  
  25.     def forward(self, x):
  26.         return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0
  27.  
  28.  
  29. class h_sigmoid(nn.Module):
  30.     def __init__(self, inplace=True, h_max=1):
  31.         super(h_sigmoid, self).__init__()
  32.         self.relu = nn.ReLU6(inplace=inplace)
  33.         self.h_max = h_max
  34.  
  35.     def forward(self, x):
  36.         return self.relu(x + 3) * self.h_max / 6
  37.  
  38.  
  39. class DYReLU(nn.Module):
  40.     def __init__(self, inp, oup, reduction=4, lambda_a=1.0, K2=True, use_bias=True, use_spatial=False,
  41.                  init_a=[1.0, 0.0], init_b=[0.0, 0.0]):
  42.         super(DYReLU, self).__init__()
  43.         self.oup = oup
  44.         self.lambda_a = lambda_a * 2
  45.         self.K2 = K2
  46.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
  47.  
  48.         self.use_bias = use_bias
  49.         if K2:
  50.             self.exp = 4 if use_bias else 2
  51.         else:
  52.             self.exp = 2 if use_bias else 1
  53.         self.init_a = init_a
  54.         self.init_b = init_b
  55.  
  56.         # determine squeeze
  57.         if reduction == 4:
  58.             squeeze = inp // reduction
  59.         else:
  60.             squeeze = _make_divisible(inp // reduction, 4)
  61.         # print('reduction: {}, squeeze: {}/{}'.format(reduction, inp, squeeze))
  62.         # print('init_a: {}, init_b: {}'.format(self.init_a, self.init_b))
  63.  
  64.         self.fc = nn.Sequential(
  65.             nn.Linear(inp, squeeze),
  66.             nn.ReLU(inplace=True),
  67.             nn.Linear(squeeze, oup * self.exp),
  68.             h_sigmoid()
  69.         )
  70.         if use_spatial:
  71.             self.spa = nn.Sequential(
  72.                 nn.Conv2d(inp, 1, kernel_size=1),
  73.                 nn.BatchNorm2d(1),
  74.             )
  75.         else:
  76.             self.spa = None
  77.  
  78.     def forward(self, x):
  79.         if isinstance(x, list):
  80.             x_in = x[0]
  81.             x_out = x[1]
  82.         else:
  83.             x_in = x
  84.             x_out = x
  85.         b, c, h, w = x_in.size()
  86.         y = self.avg_pool(x_in).view(b, c)
  87.         y = self.fc(y).view(b, self.oup * self.exp, 1, 1)
  88.         if self.exp == 4:
  89.             a1, b1, a2, b2 = torch.split(y, self.oup, dim=1)
  90.             a1 = (a1 - 0.5) * self.lambda_a + self.init_a[0]  # 1.0
  91.             a2 = (a2 - 0.5) * self.lambda_a + self.init_a[1]
  92.  
  93.             b1 = b1 - 0.5 + self.init_b[0]
  94.             b2 = b2 - 0.5 + self.init_b[1]
  95.             out = torch.max(x_out * a1 + b1, x_out * a2 + b2)
  96.         elif self.exp == 2:
  97.             if self.use_bias:  # bias but not PL
  98.                 a1, b1 = torch.split(y, self.oup, dim=1)
  99.                 a1 = (a1 - 0.5) * self.lambda_a + self.init_a[0]  # 1.0
  100.                 b1 = b1 - 0.5 + self.init_b[0]
  101.                 out = x_out * a1 + b1
  102.  
  103.             else:
  104.                 a1, a2 = torch.split(y, self.oup, dim=1)
  105.                 a1 = (a1 - 0.5) * self.lambda_a + self.init_a[0]  # 1.0
  106.                 a2 = (a2 - 0.5) * self.lambda_a + self.init_a[1]
  107.                 out = torch.max(x_out * a1, x_out * a2)
  108.  
  109.         elif self.exp == 1:
  110.             a1 = y
  111.             a1 = (a1 - 0.5) * self.lambda_a + self.init_a[0]  # 1.0
  112.             out = x_out * a1
  113.  
  114.         if self.spa:
  115.             ys = self.spa(x_in).view(b, -1)
  116.             ys = F.softmax(ys, dim=1).view(b, 1, h, w) * h * w
  117.             ys = F.hardtanh(ys, 0, 3, inplace=True) / 3
  118.             out = out * ys
  119.  
  120.         return out
  121.  
  122.  
  123. class Conv3x3Norm(torch.nn.Module):
  124.     def __init__(self, in_channels, out_channels, stride):
  125.         super(Conv3x3Norm, self).__init__()
  126.  
  127.         self.conv = ModulatedDeformConv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1)
  128.         self.bn = nn.GroupNorm(num_groups=16, num_channels=out_channels)
  129.  
  130.     def forward(self, input, **kwargs):
  131.         x = self.conv(input.contiguous(), **kwargs)
  132.         x = self.bn(x)
  133.         return x
  134.  
  135.  
  136. class DyConv(nn.Module):
  137.     def __init__(self, in_channels=256, out_channels=256, conv_func=Conv3x3Norm):
  138.         super(DyConv, self).__init__()
  139.  
  140.         self.DyConv = nn.ModuleList()
  141.         self.DyConv.append(conv_func(in_channels, out_channels, 1))
  142.         self.DyConv.append(conv_func(in_channels, out_channels, 1))
  143.         self.DyConv.append(conv_func(in_channels, out_channels, 2))
  144.  
  145.         self.AttnConv = nn.Sequential(
  146.             nn.AdaptiveAvgPool2d(1),
  147.             nn.Conv2d(in_channels, 1, kernel_size=1),
  148.             nn.ReLU(inplace=True))
  149.  
  150.         self.h_sigmoid = h_sigmoid()
  151.         self.relu = DYReLU(in_channels, out_channels)
  152.         self.offset = nn.Conv2d(in_channels, 27, kernel_size=3, stride=1, padding=1)
  153.         self.init_weights()
  154.  
  155.     def init_weights(self):
  156.         for m in self.DyConv.modules():
  157.             if isinstance(m, nn.Conv2d):
  158.                 nn.init.normal_(m.weight.data, 0, 0.01)
  159.                 if m.bias is not None:
  160.                     m.bias.data.zero_()
  161.         for m in self.AttnConv.modules():
  162.             if isinstance(m, nn.Conv2d):
  163.                 nn.init.normal_(m.weight.data, 0, 0.01)
  164.                 if m.bias is not None:
  165.                     m.bias.data.zero_()
  166.  
  167.     def forward(self, x):
  168.         next_x = {}
  169.         feature_names = list(x.keys())
  170.         for level, name in enumerate(feature_names):
  171.  
  172.             feature = x[name]
  173.  
  174.             offset_mask = self.offset(feature)
  175.             offset = offset_mask[:, :18, :, :]
  176.             mask = offset_mask[:, 18:, :, :].sigmoid()
  177.             conv_args = dict(offset=offset, mask=mask)
  178.  
  179.             temp_fea = [self.DyConv[1](feature, **conv_args)]
  180.             if level > 0:
  181.                 temp_fea.append(self.DyConv[2](x[feature_names[level - 1]], **conv_args))
  182.             if level < len(x) - 1:
  183.                 input = x[feature_names[level + 1]]
  184.                 temp_fea.append(F.interpolate(self.DyConv[0](input, **conv_args),
  185.                                               size=[feature.size(2), feature.size(3)]))
  186.             attn_fea = []
  187.             res_fea = []
  188.             for fea in temp_fea:
  189.                 res_fea.append(fea)
  190.                 attn_fea.append(self.AttnConv(fea))
  191.  
  192.             res_fea = torch.stack(res_fea)
  193.             spa_pyr_attn = self.h_sigmoid(torch.stack(attn_fea))
  194.             mean_fea = torch.mean(res_fea * spa_pyr_attn, dim=0, keepdim=False)
  195.             next_x[name] = self.relu(mean_fea)
  196.  
  197.         return next_x
  198.  
  199.  
  200. def autopad(k, p=None, d=1):  # kernel, padding, dilation
  201.     """Pad to 'same' shape outputs."""
  202.     if d > 1:
  203.         k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
  204.     if p is None:
  205.         p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
  206.     return p
  207.  
  208.  
  209. class Conv(nn.Module):
  210.     """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
  211.     default_act = nn.SiLU()  # default activation
  212.  
  213.     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
  214.         """Initialize Conv layer with given arguments including activation."""
  215.         super().__init__()
  216.         self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
  217.         self.bn = nn.BatchNorm2d(c2)
  218.         self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
  219.  
  220.     def forward(self, x):
  221.         """Apply convolution, batch normalization and activation to input tensor."""
  222.         return self.act(self.bn(self.conv(x)))
  223.  
  224.     def forward_fuse(self, x):
  225.         """Perform transposed convolution of 2D data."""
  226.         return self.act(self.conv(x))
  227.  
  228. class DFL(nn.Module):
  229.     """
  230.     Integral module of Distribution Focal Loss (DFL).
  231.     Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
  232.     """
  233.  
  234.     def __init__(self, c1=16):
  235.         """Initialize a convolutional layer with a given number of input channels."""
  236.         super().__init__()
  237.         self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
  238.         x = torch.arange(c1, dtype=torch.float)
  239.         self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
  240.         self.c1 = c1
  241.  
  242.     def forward(self, x):
  243.         """Applies a transformer layer on input tensor 'x' and returns a tensor."""
  244.         b, c, a = x.shape  # batch, channels, anchors
  245.         return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
  246.         # return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
  247.  
  248. class Detect_DynamicHead(nn.Module):
  249.     """YOLOv8 Detect head for detection models."""
  250.     dynamic = False  # force grid reconstruction
  251.     export = False  # export mode
  252.     shape = None
  253.     anchors = torch.empty(0)  # init
  254.     strides = torch.empty(0)  # init
  255.  
  256.     def __init__(self, nc=80, ch=()):
  257.         """Initializes the YOLOv8 detection layer with specified number of classes and channels."""
  258.         super().__init__()
  259.         self.nc = nc  # number of classes
  260.         self.nl = len(ch)  # number of detection layers
  261.         self.reg_max = 16  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
  262.         self.no = nc + self.reg_max * 4  # number of outputs per anchor
  263.         self.stride = torch.zeros(self.nl)  # strides computed during build
  264.         c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))  # channels
  265.         self.cv2 = nn.ModuleList(
  266.             nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
  267.         self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
  268.         self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
  269.         dyhead_tower = []
  270.         for i in range(self.nl):
  271.             channel = ch[i]
  272.             dyhead_tower.append(
  273.                 DyConv(
  274.                     channel,
  275.                     channel,
  276.                     conv_func=Conv3x3Norm,
  277.                 )
  278.             )
  279.         self.add_module('dyhead_tower', nn.Sequential(*dyhead_tower))
  280.  
  281.     def forward(self, x):
  282.         tensor_dict = {i: tensor for i, tensor in enumerate(x)}
  283.         x = self.dyhead_tower(tensor_dict)
  284.         x = list(x.values())
  285.         """Concatenates and returns predicted bounding boxes and class probabilities."""
  286.         shape = x[0].shape  # BCHW
  287.         for i in range(self.nl):
  288.             x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
  289.         if self.training:
  290.             return x
  291.         elif self.dynamic or self.shape != shape:
  292.             self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
  293.             self.shape = shape
  294.  
  295.         x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
  296.         if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'):  # avoid TF FlexSplitV ops
  297.             box = x_cat[:, :self.reg_max * 4]
  298.             cls = x_cat[:, self.reg_max * 4:]
  299.         else:
  300.             box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
  301.         dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
  302.  
  303.         if self.export and self.format in ('tflite', 'edgetpu'):
  304.             # Normalize xywh with image size to mitigate quantization error of TFLite integer models as done in YOLOv5:
  305.             # https://github.com/ultralytics/yolov5/blob/0c8de3fca4a702f8ff5c435e67f378d1fce70243/models/tf.py#L307-L309
  306.             # See this PR for details: https://github.com/ultralytics/ultralytics/pull/1695
  307.             img_h = shape[2] * self.stride[0]
  308.             img_w = shape[3] * self.stride[0]
  309.             img_size = torch.tensor([img_w, img_h, img_w, img_h], device=dbox.device).reshape(1, 4, 1)
  310.             dbox /= img_size
  311.  
  312.         y = torch.cat((dbox, cls.sigmoid()), 1)
  313.         return y if self.export else (y, x)
  314.  
  315.     def bias_init(self):
  316.         """Initialize Detect() biases, WARNING: requires stride availability."""
  317.         m = self  # self.model[-1]  # Detect() module
  318.         # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
  319.         # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
  320.         for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
  321.             a[-1].bias.data[:] = 1.0  # box
  322.             b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)

② 修改ultralytics/nn/tasks.py文件

具体的修改内容如下图所示：

可以Ctrl+f然后搜索 “detect”方便后续锁定需要修改的位置。

2 配置文件

yolov8_DynamicHead.yaml的内容如下所示：

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

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