学习资源站

YOLOv11改进-检测头篇-2024独家创新自适应性DWConv改进v11检测头独创FADWCHead（全网独家首发创新）

一、本文介绍

本文给大家带来的最新改进是独家创新利用 Frequency-Adaptive Dilated Convolution 改进YOLOv11的检测头， 频率自适应膨胀卷积（FADC）， FADC的核心思想是 根据图像的局部频率成分动态调整膨胀率 。这种方法使得网络能够根据图像内容的局部变化来调整 感受野 ，从而在细节丰富或高频信息密集的区域提高 性能 ， 本文内容为博主全网独创新 ，下图为精度对比表现。

二、原理介绍

官方论文地址： 官方论文地址点击此处即可跳转

官方代码地址： 官方代码地址点击此处即可跳转

频率自适应膨胀卷积（FADC） ，其目的是提高 膨胀卷积 在语义分割中的性能。以下是主要思想的总结：

  1. 膨胀卷积概述 ：膨胀卷积通过在卷积核的连续元素之间插入间隔，来扩展感受野。这种方法可以在不增加参数数量的情况下，捕获更广泛的上下文信息。

  2. 固定膨胀率的问题 ：传统的膨胀卷积方法将膨胀率作为全局超参数固定，而这种固定的膨胀率可能在处理图像的不同区域时不够理想，因为图像的不同部分包含不同的频率成分。

  3. 频率自适应膨胀卷积（FADC） ：FADC的核心思想是 根据图像的局部频率成分动态调整膨胀率 。这种方法使得网络能够根据图像内容的局部变化来调整感受野，从而在细节丰富或高频信息密集的区域提高性能。

  4. 两个插件模块 ：

    • 自适应卷积核（AdaKern） ：该模块通过将卷积权重分解为低频和高频部分，来增强带宽和感受野大小的调节能力。
    • 第二个模块（在文中未完全描述）进一步增强了模型对不同频率成分的适应能力。

这些方法的目标是在 语义分割 任务中，通过让卷积网络能够更好地捕捉图像特征的局部变化，来提升性能。

三、核心代码

YOLOv11的检测头使用了DWConv，这个论文提出了一种频率选择深度卷积整好可以替换，使用方式看章节四！ 

  1. import copy
  2. import math
  3. import numpy as np
  4. import torch
  5. import torch.nn as nn
  6. import torch.fft
  7. from scipy.spatial import distance
  8. from mmcv.ops.modulated_deform_conv import ModulatedDeformConv2d, modulated_deform_conv2d
  9. import torch_dct as dct
  10. from ultralytics.utils.tal import dist2bbox, make_anchors
  13. class OmniAttention(nn.Module):
  14.     def __init__(self, in_planes, out_planes, kernel_size, groups=1, reduction=0.0625, kernel_num=4, min_channel=16):
  15.         super(OmniAttention, self).__init__()
  16.         attention_channel = max(int(in_planes * reduction), min_channel)
  17.         self.kernel_size = kernel_size
  18.         self.kernel_num = kernel_num
  19.         self.temperature = 1.0
  21.         self.avgpool = nn.AdaptiveAvgPool2d(1)
  22.         self.fc = nn.Conv2d(in_planes, attention_channel, 1, bias=False)
  23.         self.bn = nn.BatchNorm2d(attention_channel)
  24.         self.relu = nn.ReLU(inplace=True)
  26.         self.channel_fc = nn.Conv2d(attention_channel, in_planes, 1, bias=True)
  27.         self.func_channel = self.get_channel_attention
  29.         if in_planes == groups and in_planes == out_planes:  # depth-wise convolution
  30.             self.func_filter = self.skip
  31.         else:
  32.             self.filter_fc = nn.Conv2d(attention_channel, out_planes, 1, bias=True)
  33.             self.func_filter = self.get_filter_attention
  35.         if kernel_size == 1:  # point-wise convolution
  36.             self.func_spatial = self.skip
  37.         else:
  38.             self.spatial_fc = nn.Conv2d(attention_channel, kernel_size * kernel_size, 1, bias=True)
  39.             self.func_spatial = self.get_spatial_attention
  41.         if kernel_num == 1:
  42.             self.func_kernel = self.skip
  43.         else:
  44.             self.kernel_fc = nn.Conv2d(attention_channel, kernel_num, 1, bias=True)
  45.             self.func_kernel = self.get_kernel_attention
  47.         self._initialize_weights()
  49.     def _initialize_weights(self):
  50.         for m in self.modules():
  51.             if isinstance(m, nn.Conv2d):
  52.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  53.                 if m.bias is not None:
  54.                     nn.init.constant_(m.bias, 0)
  55.             if isinstance(m, nn.BatchNorm2d):
  56.                 nn.init.constant_(m.weight, 1)
  57.                 nn.init.constant_(m.bias, 0)
  59.     def update_temperature(self, temperature):
  60.         self.temperature = temperature
  62.     @staticmethod
  63.     def skip(_):
  64.         return 1.0
  66.     def get_channel_attention(self, x):
  67.         channel_attention = torch.sigmoid(self.channel_fc(x).view(x.size(0), -1, 1, 1) / self.temperature)
  68.         return channel_attention
  70.     def get_filter_attention(self, x):
  71.         filter_attention = torch.sigmoid(self.filter_fc(x).view(x.size(0), -1, 1, 1) / self.temperature)
  72.         return filter_attention
  74.     def get_spatial_attention(self, x):
  75.         spatial_attention = self.spatial_fc(x).view(x.size(0), 1, 1, 1, self.kernel_size, self.kernel_size)
  76.         spatial_attention = torch.sigmoid(spatial_attention / self.temperature)
  77.         return spatial_attention
  79.     def get_kernel_attention(self, x):
  80.         kernel_attention = self.kernel_fc(x).view(x.size(0), -1, 1, 1, 1, 1)
  81.         kernel_attention = F.softmax(kernel_attention / self.temperature, dim=1)
  82.         return kernel_attention
  84.     def forward(self, x):
  85.         x = self.avgpool(x)
  86.         x = self.fc(x)
  87.         if x.shape[3] == 1:
  88.             pass
  89.         else:
  90.             x = self.bn(x)
  91.         x = self.relu(x)
  92.         return self.func_channel(x), self.func_filter(x), self.func_spatial(x), self.func_kernel(x)
  95. import torch.nn.functional as F
  98. def generate_laplacian_pyramid(input_tensor, num_levels, size_align=True, mode='bilinear'):
  99.     pyramid = []
  100.     current_tensor = input_tensor
  101.     _, _, H, W = current_tensor.shape
  102.     for _ in range(num_levels):
  103.         b, _, h, w = current_tensor.shape
  104.         downsampled_tensor = F.interpolate(current_tensor, (h // 2 + h % 2, w // 2 + w % 2), mode=mode,
  105.                                            align_corners=(H % 2) == 1)  # antialias=True
  106.         if size_align:
  107.             # upsampled_tensor = F.interpolate(downsampled_tensor, (h, w), mode='bilinear', align_corners=(H%2) == 1)
  108.             # laplacian = current_tensor - upsampled_tensor
  109.             # laplacian = F.interpolate(laplacian, (H, W), mode='bilinear', align_corners=(H%2) == 1)
  110.             upsampled_tensor = F.interpolate(downsampled_tensor, (H, W), mode=mode, align_corners=(H % 2) == 1)
  111.             laplacian = F.interpolate(current_tensor, (H, W), mode=mode, align_corners=(H % 2) == 1) - upsampled_tensor
  112.             # print(laplacian.shape)
  113.         else:
  114.             upsampled_tensor = F.interpolate(downsampled_tensor, (h, w), mode=mode, align_corners=(H % 2) == 1)
  115.             laplacian = current_tensor - upsampled_tensor
  116.         pyramid.append(laplacian)
  117.         current_tensor = downsampled_tensor
  118.     if size_align: current_tensor = F.interpolate(current_tensor, (H, W), mode=mode, align_corners=(H % 2) == 1)
  119.     pyramid.append(current_tensor)
  120.     return pyramid
  123. class FrequencySelection(nn.Module):
  124.     def __init__(self,
  125.                  in_channels,
  126.                  k_list=[2],
  127.                  # freq_list=[2, 3, 5, 7, 9, 11],
  128.                  lowfreq_att=True,
  129.                  fs_feat='feat',
  130.                  lp_type='freq',
  131.                  act='sigmoid',
  132.                  spatial='conv',
  133.                  spatial_group=1,
  134.                  spatial_kernel=3,
  135.                  init='zero',
  136.                  global_selection=False,
  137.                  ):
  138.         super().__init__()
  139.         # k_list.sort()
  140.         # print()
  141.         self.k_list = k_list
  142.         # self.freq_list = freq_list
  143.         self.lp_list = nn.ModuleList()
  144.         self.freq_weight_conv_list = nn.ModuleList()
  145.         self.fs_feat = fs_feat
  146.         self.lp_type = lp_type
  147.         self.in_channels = in_channels
  148.         # self.residual = residual
  149.         if spatial_group > 64: spatial_group = in_channels
  150.         self.spatial_group = spatial_group
  151.         self.lowfreq_att = lowfreq_att
  152.         if spatial == 'conv':
  153.             self.freq_weight_conv_list = nn.ModuleList()
  154.             _n = len(k_list)
  155.             if lowfreq_att:  _n += 1
  156.             for i in range(_n):
  157.                 freq_weight_conv = nn.Conv2d(in_channels=in_channels,
  158.                                              out_channels=self.spatial_group,
  159.                                              stride=1,
  160.                                              kernel_size=spatial_kernel,
  161.                                              groups=self.spatial_group,
  162.                                              padding=spatial_kernel // 2,
  163.                                              bias=True)
  164.                 if init == 'zero':
  165.                     freq_weight_conv.weight.data.zero_()
  166.                     freq_weight_conv.bias.data.zero_()
  167.                 else:
  168.                     # raise NotImplementedError
  169.                     pass
  170.                 self.freq_weight_conv_list.append(freq_weight_conv)
  171.         else:
  172.             raise NotImplementedError
  174.         if self.lp_type == 'avgpool':
  175.             for k in k_list:
  176.                 self.lp_list.append(nn.Sequential(
  177.                     nn.ReplicationPad2d(padding=k // 2),
  178.                     # nn.ZeroPad2d(padding= k // 2),
  179.                     nn.AvgPool2d(kernel_size=k, padding=0, stride=1)
  180.                 ))
  181.         elif self.lp_type == 'laplacian':
  182.             pass
  183.         elif self.lp_type == 'freq':
  184.             pass
  185.         else:
  186.             raise NotImplementedError
  188.         self.act = act
  189.         # self.freq_weight_conv_list.append(nn.Conv2d(self.deform_groups * 3 * self.kernel_size[0] * self.kernel_size[1], 1, kernel_size=1, padding=0, bias=True))
  190.         self.global_selection = global_selection
  191.         if self.global_selection:
  192.             self.global_selection_conv_real = nn.Conv2d(in_channels=in_channels,
  193.                                                         out_channels=self.spatial_group,
  194.                                                         stride=1,
  195.                                                         kernel_size=1,
  196.                                                         groups=self.spatial_group,
  197.                                                         padding=0,
  198.                                                         bias=True)
  199.             self.global_selection_conv_imag = nn.Conv2d(in_channels=in_channels,
  200.                                                         out_channels=self.spatial_group,
  201.                                                         stride=1,
  202.                                                         kernel_size=1,
  203.                                                         groups=self.spatial_group,
  204.                                                         padding=0,
  205.                                                         bias=True)
  206.             if init == 'zero':
  207.                 self.global_selection_conv_real.weight.data.zero_()
  208.                 self.global_selection_conv_real.bias.data.zero_()
  209.                 self.global_selection_conv_imag.weight.data.zero_()
  210.                 self.global_selection_conv_imag.bias.data.zero_()
  212.     def sp_act(self, freq_weight):
  213.         if self.act == 'sigmoid':
  214.             freq_weight = freq_weight.sigmoid() * 2
  215.         elif self.act == 'softmax':
  216.             freq_weight = freq_weight.softmax(dim=1) * freq_weight.shape[1]
  217.         else:
  218.             raise NotImplementedError
  219.         return freq_weight
  221.     def forward(self, x, att_feat=None):
  222.         """
  223.         att_feat:feat for gen att
  224.         """
  225.         # freq_weight = self.freq_weight_conv(x)
  226.         # self.sp_act(freq_weight)
  227.         # if self.residual: x_residual = x.clone()
  228.         if att_feat is None: att_feat = x
  229.         x_list = []
  230.         if self.lp_type == 'avgpool':
  231.             # for avg, freq_weight in zip(self.avg_list, self.freq_weight_conv_list):
  232.             pre_x = x
  233.             b, _, h, w = x.shape
  234.             for idx, avg in enumerate(self.lp_list):
  235.                 low_part = avg(x)
  236.                 high_part = pre_x - low_part
  237.                 pre_x = low_part
  238.                 # x_list.append(freq_weight[:, idx:idx+1] * high_part)
  239.                 freq_weight = self.freq_weight_conv_list[idx](att_feat)
  240.                 freq_weight = self.sp_act(freq_weight)
  241.                 # tmp = freq_weight[:, :, idx:idx+1] * high_part.reshape(b, self.spatial_group, -1, h, w)
  242.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * high_part.reshape(b, self.spatial_group,
  243.                                                                                                -1, h, w)
  244.                 x_list.append(tmp.reshape(b, -1, h, w))
  245.             if self.lowfreq_att:
  246.                 freq_weight = self.freq_weight_conv_list[len(x_list)](att_feat)
  247.                 # tmp = freq_weight[:, :, len(x_list):len(x_list)+1] * pre_x.reshape(b, self.spatial_group, -1, h, w)
  248.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * pre_x.reshape(b, self.spatial_group, -1, h,
  249.                                                                                            w)
  250.                 x_list.append(tmp.reshape(b, -1, h, w))
  251.             else:
  252.                 x_list.append(pre_x)
  253.         elif self.lp_type == 'laplacian':
  254.             # for avg, freq_weight in zip(self.avg_list, self.freq_weight_conv_list):
  255.             # pre_x = x
  256.             b, _, h, w = x.shape
  257.             pyramids = generate_laplacian_pyramid(x, len(self.k_list), size_align=True)
  258.             # print('pyramids', len(pyramids))
  259.             for idx, avg in enumerate(self.k_list):
  260.                 # print(idx)
  261.                 high_part = pyramids[idx]
  262.                 freq_weight = self.freq_weight_conv_list[idx](att_feat)
  263.                 freq_weight = self.sp_act(freq_weight)
  264.                 # tmp = freq_weight[:, :, idx:idx+1] * high_part.reshape(b, self.spatial_group, -1, h, w)
  265.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * high_part.reshape(b, self.spatial_group,
  266.                                                                                                -1, h, w)
  267.                 x_list.append(tmp.reshape(b, -1, h, w))
  268.             if self.lowfreq_att:
  269.                 freq_weight = self.freq_weight_conv_list[len(x_list)](att_feat)
  270.                 # tmp = freq_weight[:, :, len(x_list):len(x_list)+1] * pre_x.reshape(b, self.spatial_group, -1, h, w)
  271.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * pyramids[-1].reshape(b, self.spatial_group,
  272.                                                                                                   -1, h, w)
  273.                 x_list.append(tmp.reshape(b, -1, h, w))
  274.             else:
  275.                 x_list.append(pyramids[-1])
  276.         elif self.lp_type == 'freq':
  277.             pre_x = x.clone()
  278.             b, _, h, w = x.shape
  279.             # b, _c, h, w = freq_weight.shape
  280.             # freq_weight = freq_weight.reshape(b, self.spatial_group, -1, h, w)
  281.             x_fft = torch.fft.fftshift(torch.fft.fft2(x, norm='ortho'))
  282.             if self.global_selection:
  283.                 # global_att_real = self.global_selection_conv_real(x_fft.real)
  284.                 # global_att_real = self.sp_act(global_att_real).reshape(b, self.spatial_group, -1, h, w)
  285.                 # global_att_imag = self.global_selection_conv_imag(x_fft.imag)
  286.                 # global_att_imag = self.sp_act(global_att_imag).reshape(b, self.spatial_group, -1, h, w)
  287.                 # x_fft = x_fft.reshape(b, self.spatial_group, -1, h, w)
  288.                 # x_fft.real *= global_att_real
  289.                 # x_fft.imag *= global_att_imag
  290.                 # x_fft = x_fft.reshape(b, -1, h, w)
  291.                 # 将x_fft复数拆分成实部和虚部
  292.                 x_real = x_fft.real
  293.                 x_imag = x_fft.imag
  294.                 # 计算实部的全局注意力
  295.                 global_att_real = self.global_selection_conv_real(x_real)
  296.                 global_att_real = self.sp_act(global_att_real).reshape(b, self.spatial_group, -1, h, w)
  297.                 # 计算虚部的全局注意力
  298.                 global_att_imag = self.global_selection_conv_imag(x_imag)
  299.                 global_att_imag = self.sp_act(global_att_imag).reshape(b, self.spatial_group, -1, h, w)
  300.                 # 重塑x_fft为形状为(b, self.spatial_group, -1, h, w)的张量
  301.                 x_real = x_real.reshape(b, self.spatial_group, -1, h, w)
  302.                 x_imag = x_imag.reshape(b, self.spatial_group, -1, h, w)
  303.                 # 分别应用实部和虚部的全局注意力
  304.                 x_fft_real_updated = x_real * global_att_real
  305.                 x_fft_imag_updated = x_imag * global_att_imag
  306.                 # 合并为复数
  307.                 x_fft_updated = torch.complex(x_fft_real_updated, x_fft_imag_updated)
  308.                 # 重塑x_fft为形状为(b, -1, h, w)的张量
  309.                 x_fft = x_fft_updated.reshape(b, -1, h, w)
  311.             for idx, freq in enumerate(self.k_list):
  312.                 mask = torch.zeros_like(x[:, 0:1, :, :], device=x.device)
  313.                 mask[:, :, round(h / 2 - h / (2 * freq)):round(h / 2 + h / (2 * freq)),
  314.                 round(w / 2 - w / (2 * freq)):round(w / 2 + w / (2 * freq))] = 1.0
  315.                 low_part = torch.fft.ifft2(torch.fft.ifftshift(x_fft * mask), norm='ortho').real
  316.                 high_part = pre_x - low_part
  317.                 pre_x = low_part
  318.                 freq_weight = self.freq_weight_conv_list[idx](att_feat)
  319.                 freq_weight = self.sp_act(freq_weight)
  320.                 # tmp = freq_weight[:, :, idx:idx+1] * high_part.reshape(b, self.spatial_group, -1, h, w)
  321.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * high_part.reshape(b, self.spatial_group,
  322.                                                                                                -1, h, w)
  323.                 x_list.append(tmp.reshape(b, -1, h, w))
  324.             if self.lowfreq_att:
  325.                 freq_weight = self.freq_weight_conv_list[len(x_list)](att_feat)
  326.                 # tmp = freq_weight[:, :, len(x_list):len(x_list)+1] * pre_x.reshape(b, self.spatial_group, -1, h, w)
  327.                 tmp = freq_weight.reshape(b, self.spatial_group, -1, h, w) * pre_x.reshape(b, self.spatial_group, -1, h,
  328.                                                                                            w)
  329.                 x_list.append(tmp.reshape(b, -1, h, w))
  330.             else:
  331.                 x_list.append(pre_x)
  332.         x = sum(x_list)
  333.         return x
  336. class AdaptiveDilatedConv(ModulatedDeformConv2d):
  337.     """A ModulatedDeformable Conv Encapsulation that acts as normal Conv
  338.     layers.
  340.     Args:
  341.         in_channels (int): Same as nn.Conv2d.
  342.         out_channels (int): Same as nn.Conv2d.
  343.         kernel_size (int or tuple[int]): Same as nn.Conv2d.
  344.         stride (int): Same as nn.Conv2d, while tuple is not supported.
  345.         padding (int): Same as nn.Conv2d, while tuple is not supported.
  346.         dilation (int): Same as nn.Conv2d, while tuple is not supported.
  347.         groups (int): Same as nn.Conv2d.
  348.         bias (bool or str): If specified as `auto`, it will be decided by the
  349.             norm_cfg. Bias will be set as True if norm_cfg is None, otherwise
  350.             False.
  351.     """
  353.     _version = 2
  355.     def __init__(self, *args,
  356.                  offset_freq=None,  # deprecated
  357.                  padding_mode='repeat',
  358.                  kernel_decompose='both',
  359.                  conv_type='conv',
  360.                  sp_att=False,
  361.                  pre_fs=True,  # False, use dilation
  362.                  epsilon=1e-4,
  363.                  use_zero_dilation=False,
  364.                  use_dct=False,
  365.                  fs_cfg={
  366.                      'k_list': [2, 4, 8],
  367.                      'fs_feat': 'feat',
  368.                      'lowfreq_att': False,
  369.                      'lp_type': 'freq',
  370.                      # 'lp_type':'laplacian',
  371.                      'act': 'sigmoid',
  372.                      'spatial': 'conv',
  373.                      'spatial_group': 1,
  374.                  },
  375.                  **kwargs):
  376.         super().__init__(*args, **kwargs)
  377.         if padding_mode == 'zero':
  378.             self.PAD = nn.ZeroPad2d(self.kernel_size[0] // 2)
  379.         elif padding_mode == 'repeat':
  380.             self.PAD = nn.ReplicationPad2d(self.kernel_size[0] // 2)
  381.         else:
  382.             self.PAD = nn.Identity()
  384.         self.kernel_decompose = kernel_decompose
  385.         self.use_dct = use_dct
  387.         if kernel_decompose == 'both':
  388.             self.OMNI_ATT1 = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  389.                                            groups=1, reduction=0.0625, kernel_num=1, min_channel=16)
  390.             self.OMNI_ATT2 = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels,
  391.                                            kernel_size=self.kernel_size[0] if self.use_dct else 1, groups=1,
  392.                                            reduction=0.0625, kernel_num=1, min_channel=16)
  393.         elif kernel_decompose == 'high':
  394.             self.OMNI_ATT = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  395.                                           groups=1, reduction=0.0625, kernel_num=1, min_channel=16)
  396.         elif kernel_decompose == 'low':
  397.             self.OMNI_ATT = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  398.                                           groups=1, reduction=0.0625, kernel_num=1, min_channel=16)
  399.         self.conv_type = conv_type
  400.         if conv_type == 'conv':
  401.             self.conv_offset = nn.Conv2d(
  402.                 self.in_channels,
  403.                 self.deform_groups * 1,
  404.                 kernel_size=self.kernel_size,
  405.                 stride=self.stride,
  406.                 padding=self.kernel_size[0] // 2 if isinstance(self.PAD, nn.Identity) else 0,
  407.                 dilation=1,
  408.                 bias=True)
  409.         else:
  410.             raise NotImplementedError
  411.             pass
  412.         # self.conv_offset_low = nn.Sequential(
  413.         #     nn.AvgPool2d(
  414.         #         kernel_size=self.kernel_size,
  415.         #         stride=self.stride,
  416.         #         padding=1,
  417.         #     ),
  418.         #     nn.Conv2d(
  419.         #         self.in_channels,
  420.         #         self.deform_groups * 1,
  421.         #         kernel_size=1,
  422.         #         stride=1,
  423.         #         padding=0,
  424.         #         dilation=1,
  425.         #         bias=False),
  426.         # )
  428.         # self.conv_offset_high = nn.Sequential(
  429.         #     LHPFConv3(channels=self.in_channels, stride=1, padding=1, residual=False),
  430.         #     nn.Conv2d(
  431.         #         self.in_channels,
  432.         #         self.deform_groups * 1,
  433.         #         kernel_size=1,
  434.         #         stride=1,
  435.         #         padding=0,
  436.         #         dilation=1,
  437.         #         bias=True),
  438.         # )
  439.         self.conv_mask = nn.Conv2d(
  440.             self.in_channels,
  441.             self.deform_groups * 1 * self.kernel_size[0] * self.kernel_size[1],
  442.             kernel_size=self.kernel_size,
  443.             stride=self.stride,
  444.             padding=self.kernel_size[0] // 2 if isinstance(self.PAD, nn.Identity) else 0,
  445.             dilation=1,
  446.             bias=True)
  447.         if sp_att:
  448.             self.conv_mask_mean_level = nn.Conv2d(
  449.                 self.in_channels,
  450.                 self.deform_groups * 1,
  451.                 kernel_size=self.kernel_size,
  452.                 stride=self.stride,
  453.                 padding=self.kernel_size[0] // 2 if isinstance(self.PAD, nn.Identity) else 0,
  454.                 dilation=1,
  455.                 bias=True)
  457.         self.offset_freq = offset_freq
  460.         # An offset is like [y0, x0, y1, x1, y2, x2, ⋯, y8, x8]
  461.         offset = [-1, -1, -1, 0, -1, 1,
  462.                   0, -1, 0, 0, 0, 1,
  463.                   1, -1, 1, 0, 1, 1]
  464.         offset = torch.Tensor(offset)
  465.         # offset[0::2] *= self.dilation[0]
  466.         # offset[1::2] *= self.dilation[1]
  467.         # a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension
  468.         self.register_buffer('dilated_offset', torch.Tensor(offset[None, None, ..., None, None]))  # B, G, 18, 1, 1
  469.         if fs_cfg is not None:
  470.             if pre_fs:
  471.                 self.FS = FrequencySelection(self.in_channels, **fs_cfg)
  472.             else:
  473.                 self.FS = FrequencySelection(1, **fs_cfg)  # use dilation
  474.         self.pre_fs = pre_fs
  475.         self.epsilon = epsilon
  476.         self.use_zero_dilation = use_zero_dilation
  477.         self.init_weights()
  479.     def freq_select(self, x):
  480.         if self.offset_freq is None:
  481.             res = x
  482.         elif self.offset_freq in ('FLC_high', 'SLP_high'):
  483.             res = x - self.LP(x)
  484.         elif self.offset_freq in ('FLC_res', 'SLP_res'):
  485.             res = 2 * x - self.LP(x)
  486.         else:
  487.             raise NotImplementedError
  488.         return res
  490.     def init_weights(self):
  491.         super().init_weights()
  492.         if hasattr(self, 'conv_offset'):
  493.             # if isinstanace(self.conv_offset, nn.Conv2d):
  494.             if self.conv_type == 'conv':
  495.                 self.conv_offset.weight.data.zero_()
  496.                 # self.conv_offset.bias.data.fill_((self.dilation[0] - 1) / self.dilation[0] + 1e-4)
  497.                 self.conv_offset.bias.data.fill_((self.dilation[0] - 1) / self.dilation[0] + self.epsilon)
  498.             # self.conv_offset.bias.data.zero_()
  499.         if hasattr(self, 'conv_mask'):
  500.             self.conv_mask.weight.data.zero_()
  501.             self.conv_mask.bias.data.zero_()
  503.         if hasattr(self, 'conv_mask_mean_level'):
  504.             self.conv_mask.weight.data.zero_()
  505.             self.conv_mask.bias.data.zero_()
  507.     # @force_fp32(apply_to=('x',))
  508.     # @force_fp32
  509.     def forward(self, x):
  510.         # offset = self.conv_offset(self.freq_select(x)) + self.conv_offset_low(self.freq_select(x))
  511.         if hasattr(self, 'FS') and self.pre_fs: x = self.FS(x)
  512.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  513.             c_att1, f_att1, _, _, = self.OMNI_ATT1(x)
  514.             c_att2, f_att2, spatial_att2, _, = self.OMNI_ATT2(x)
  515.         elif hasattr(self, 'OMNI_ATT'):
  516.             c_att, f_att, _, _, = self.OMNI_ATT(x)
  518.         if self.conv_type == 'conv':
  519.             offset = self.conv_offset(self.PAD(self.freq_select(x)))
  520.         elif self.conv_type == 'multifreqband':
  521.             offset = self.conv_offset(self.freq_select(x))
  522.         # high_gate = self.conv_offset_high(x)
  523.         # high_gate = torch.exp(-0.5 * high_gate ** 2)
  524.         # offset = F.relu(offset, inplace=True) * self.dilation[0] - 1 # ensure > 0
  525.         if self.use_zero_dilation:
  526.             offset = (F.relu(offset + 1, inplace=True) - 1) * self.dilation[0]  # ensure > 0
  527.         else:
  528.             # offset = F.relu(offset, inplace=True) * self.dilation[0] # ensure > 0
  529.             offset = offset.abs() * self.dilation[0]  # ensure > 0
  530.             # offset[offset<0] = offset[offset<0].exp() - 1
  531.         # print(offset.mean(), offset.std(), offset.max(), offset.min())
  532.         if hasattr(self, 'FS') and (self.pre_fs == False): x = self.FS(x, F.interpolate(offset, x.shape[-2:],
  533.                                                                                         mode='bilinear', align_corners=(
  534.                                                                                                                                    x.shape[
  535.                                                                                                                                        -1] % 2) == 1))
  536.         # print(offset.max(), offset.abs().min(), offset.abs().mean())
  537.         # offset *= high_gate # ensure > 0
  538.         b, _, h, w = offset.shape
  539.         offset = offset.reshape(b, self.deform_groups, -1, h, w) * self.dilated_offset
  540.         # offset = offset.reshape(b, self.deform_groups, -1, h, w).repeat(1, 1, 9, 1, 1)
  541.         # offset[:, :, 0::2, ] *= self.dilated_offset[:, :, 0::2, ]
  542.         # offset[:, :, 1::2, ] *= self.dilated_offset[:, :, 1::2, ]
  543.         offset = offset.reshape(b, -1, h, w)
  545.         x = self.PAD(x)
  546.         mask = self.conv_mask(x)
  547.         mask = mask.sigmoid()
  548.         # print(mask.shape)
  549.         # mask = mask.reshape(b, self.deform_groups, -1, h, w).softmax(dim=2)
  550.         if hasattr(self, 'conv_mask_mean_level'):
  551.             mask_mean_level = torch.sigmoid(self.conv_mask_mean_level(x)).reshape(b, self.deform_groups, -1, h, w)
  552.             mask = mask * mask_mean_level
  553.         mask = mask.reshape(b, -1, h, w)
  555.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  556.             offset = offset.reshape(1, -1, h, w)
  557.             mask = mask.reshape(1, -1, h, w)
  558.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  559.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, c_out, c_in, k, k
  560.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  561.             adaptive_weight_res = adaptive_weight - adaptive_weight_mean
  562.             b, c_out, c_in, k, k = adaptive_weight.shape
  563.             if self.use_dct:
  564.                 dct_coefficients = dct.dct_2d(adaptive_weight_res)
  565.                 # print(adaptive_weight_res.shape, dct_coefficients.shape)
  566.                 spatial_att2 = spatial_att2.reshape(b, 1, 1, k, k)
  567.                 dct_coefficients = dct_coefficients * (spatial_att2 * 2)
  568.                 # print(dct_coefficients.shape)
  569.                 adaptive_weight_res = dct.idct_2d(dct_coefficients)
  570.                 # adaptive_weight_res = adaptive_weight_res.reshape(b, c_out, c_in, k, k)
  571.                 # print(adaptive_weight_res.shape, dct_coefficients.shape)
  572.             # adaptive_weight = adaptive_weight_mean * (2 * c_att.unsqueeze(1)) * (2 * f_att.unsqueeze(2)) + adaptive_weight - adaptive_weight_mean
  573.             # adaptive_weight = adaptive_weight_mean * (c_att1.unsqueeze(1) * 2) * (f_att1.unsqueeze(2) * 2) + (adaptive_weight - adaptive_weight_mean) * (c_att2.unsqueeze(1) * 2) * (f_att2.unsqueeze(2) * 2)
  574.             adaptive_weight = adaptive_weight_mean * (c_att1.unsqueeze(1) * 2) * (
  575.                         f_att1.unsqueeze(2) * 2) + adaptive_weight_res * (c_att2.unsqueeze(1) * 2) * (
  576.                                           f_att2.unsqueeze(2) * 2)
  577.             adaptive_weight = adaptive_weight.reshape(-1, self.in_channels // self.groups, 3, 3)
  578.             if self.bias is not None:
  579.                 bias = self.bias.repeat(b)
  580.             else:
  581.                 bias = self.bias
  582.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight, bias,
  583.                                         self.stride,
  584.                                         (self.kernel_size[0] // 2, self.kernel_size[1] // 2) if isinstance(self.PAD,
  585.                                                                                                            nn.Identity) else (
  586.                                         0, 0),  # padding
  587.                                         (1, 1),  # dilation
  588.                                         self.groups * b, self.deform_groups * b)
  589.         elif hasattr(self, 'OMNI_ATT'):
  590.             offset = offset.reshape(1, -1, h, w)
  591.             mask = mask.reshape(1, -1, h, w)
  592.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  593.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, c_out, c_in, k, k
  594.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  595.             # adaptive_weight = adaptive_weight_mean * (2 * c_att.unsqueeze(1)) * (2 * f_att.unsqueeze(2)) + adaptive_weight - adaptive_weight_mean
  596.             if self.kernel_decompose == 'high':
  597.                 adaptive_weight = adaptive_weight_mean + (adaptive_weight - adaptive_weight_mean) * (
  598.                             c_att.unsqueeze(1) * 2) * (f_att.unsqueeze(2) * 2)
  599.             elif self.kernel_decompose == 'low':
  600.                 adaptive_weight = adaptive_weight_mean * (c_att.unsqueeze(1) * 2) * (f_att.unsqueeze(2) * 2) + (
  601.                             adaptive_weight - adaptive_weight_mean)
  603.             adaptive_weight = adaptive_weight.reshape(-1, self.in_channels // self.groups, 3, 3)
  604.             if self.bias is not None:
  605.                 bias = self.bias.repeat(b)
  606.             else:
  607.                 bias = self.bias
  608.             # adaptive_bias = self.unsqueeze(0).repeat(b, 1, 1, 1, 1)
  609.             # print(adaptive_weight.shape)
  610.             # print(offset.shape)
  611.             # print(mask.shape)
  612.             # print(x.shape)
  613.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight, bias,
  614.                                         self.stride,
  615.                                         (self.kernel_size[0] // 2, self.kernel_size[1] // 2) if isinstance(self.PAD,
  616.                                                                                                            nn.Identity) else (
  617.                                         0, 0),  # padding
  618.                                         (1, 1),  # dilation
  619.                                         self.groups * b, self.deform_groups * b)
  620.         else:
  621.             x = modulated_deform_conv2d(x, offset, mask, self.weight, self.bias,
  622.                                         self.stride,
  623.                                         (self.kernel_size[0] // 2, self.kernel_size[1] // 2) if isinstance(self.PAD,
  624.                                                                                                            nn.Identity) else (
  625.                                         0, 0),  # padding
  626.                                         (1, 1),  # dilation
  627.                                         self.groups, self.deform_groups)
  628.         # x = modulated_deform_conv2d(x, offset, mask, self.weight, self.bias,
  629.         #                                self.stride, self.padding,
  630.         #                                self.dilation, self.groups,
  631.         #                                self.deform_groups)
  632.         # if hasattr(self, 'OMNI_ATT'): x = x * f_att
  633.         return x.reshape(b, -1, h, w)
  637. class AdaptiveDilatedDWConv(ModulatedDeformConv2d):
  638.     """A ModulatedDeformable Conv Encapsulation that acts as normal Conv
  639.     layers.
  641.     Args:
  642.         in_channels (int): Same as nn.Conv2d.
  643.         out_channels (int): Same as nn.Conv2d.
  644.         kernel_size (int or tuple[int]): Same as nn.Conv2d.
  645.         stride (int): Same as nn.Conv2d, while tuple is not supported.
  646.         padding (int): Same as nn.Conv2d, while tuple is not supported.
  647.         dilation (int): Same as nn.Conv2d, while tuple is not supported.
  648.         groups (int): Same as nn.Conv2d.
  649.         bias (bool or str): If specified as `auto`, it will be decided by the
  650.             norm_cfg. Bias will be set as True if norm_cfg is None, otherwise
  651.             False.
  652.     """
  654.     _version = 2
  656.     def __init__(self, *args,
  657.                  offset_freq=None,
  658.                  use_BFM=False,
  659.                  kernel_decompose='both',
  660.                  padding_mode='repeat',
  661.                  #  padding_mode='zero',
  662.                  normal_conv_dim=0,
  663.                  pre_fs=True,  # False, use dilation
  664.                  fs_cfg={
  665.                      # 'k_list':[3,5,7,9],
  666.                      'k_list': [2, 4, 8],
  667.                      'fs_feat': 'feat',
  668.                      'lowfreq_att': False,
  669.                      # 'lp_type':'freq_eca',
  670.                      # 'lp_type':'freq_channel_att',
  671.                      # 'lp_type':'freq',
  672.                      # 'lp_type':'avgpool',
  673.                      'lp_type': 'freq',
  674.                      'act': 'sigmoid',
  675.                      'spatial': 'conv',
  676.                      'spatial_group': 1,
  677.                  },
  678.                  **kwargs):
  679.         super().__init__(*args, **kwargs)
  680.         assert self.kernel_size[0] in (3, 7)
  681.         assert self.groups == self.in_channels
  682.         if kernel_decompose == 'both':
  683.             self.OMNI_ATT1 = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  684.                                            groups=self.in_channels, reduction=0.0625, kernel_num=1, min_channel=16)
  685.             self.OMNI_ATT2 = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  686.                                            groups=self.in_channels, reduction=0.0625, kernel_num=1, min_channel=16)
  687.         elif kernel_decompose == 'high':
  688.             self.OMNI_ATT = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  689.                                           groups=self.in_channels, reduction=0.0625, kernel_num=1, min_channel=16)
  690.         elif kernel_decompose == 'low':
  691.             self.OMNI_ATT = OmniAttention(in_planes=self.in_channels, out_planes=self.out_channels, kernel_size=1,
  692.                                           groups=self.in_channels, reduction=0.0625, kernel_num=1, min_channel=16)
  693.         self.kernel_decompose = kernel_decompose
  695.         self.normal_conv_dim = normal_conv_dim
  697.         if padding_mode == 'zero':
  698.             self.PAD = nn.ZeroPad2d(self.kernel_size[0] // 2)
  699.         elif padding_mode == 'repeat':
  700.             self.PAD = nn.ReplicationPad2d(self.kernel_size[0] // 2)
  701.         else:
  702.             self.PAD = nn.Identity()
  703.         self.conv_offset = nn.Conv2d(
  704.             self.in_channels - self.normal_conv_dim,
  705.             self.deform_groups * 1,
  706.             # self.groups * 1,
  707.             kernel_size=self.kernel_size,
  708.             stride=self.stride,
  709.             padding=self.padding if isinstance(self.PAD, nn.Identity) else 0,
  710.             dilation=1,
  711.             bias=True)
  712.         # self.conv_offset_low = nn.Sequential(
  713.         #     nn.AvgPool2d(
  714.         #         kernel_size=self.kernel_size,
  715.         #         stride=self.stride,
  716.         #         padding=1,
  717.         #     ),
  718.         #     nn.Conv2d(
  719.         #         self.in_channels,
  720.         #         self.deform_groups * 1,
  721.         #         kernel_size=1,
  722.         #         stride=1,
  723.         #         padding=0,
  724.         #         dilation=1,
  725.         #         bias=False),
  726.         # )
  727.         self.conv_mask = nn.Sequential(
  728.             nn.Conv2d(
  729.                 self.in_channels - self.normal_conv_dim,
  730.                 self.in_channels - self.normal_conv_dim,
  731.                 kernel_size=self.kernel_size,
  732.                 stride=self.stride,
  733.                 padding=self.padding if isinstance(self.PAD, nn.Identity) else 0,
  734.                 groups=self.in_channels - self.normal_conv_dim,
  735.                 dilation=1,
  736.                 bias=False),
  737.             nn.Conv2d(
  738.                 self.in_channels - self.normal_conv_dim,
  739.                 self.deform_groups * 1 * self.kernel_size[0] * self.kernel_size[1],
  740.                 kernel_size=1,
  741.                 stride=1,
  742.                 padding=0,
  743.                 groups=1,
  744.                 dilation=1,
  745.                 bias=True)
  746.         )
  748.         self.offset_freq = offset_freq
  752.         # An offset is like [y0, x0, y1, x1, y2, x2, ⋯, y8, x8]
  753.         if self.kernel_size[0] == 3:
  754.             offset = [-1, -1, -1, 0, -1, 1,
  755.                       0, -1, 0, 0, 0, 1,
  756.                       1, -1, 1, 0, 1, 1]
  757.         elif self.kernel_size[0] == 7:
  758.             offset = [
  759.                 -3, -3, -3, -2, -3, -1, -3, 0, -3, 1, -3, 2, -3, 3,
  760.                 -2, -3, -2, -2, -2, -1, -2, 0, -2, 1, -2, 2, -2, 3,
  761.                 -1, -3, -1, -2, -1, -1, -1, 0, -1, 1, -1, 2, -1, 3,
  762.                 0, -3, 0, -2, 0, -1, 0, 0, 0, 1, 0, 2, 0, 3,
  763.                 1, -3, 1, -2, 1, -1, 1, 0, 1, 1, 1, 2, 1, 3,
  764.                 2, -3, 2, -2, 2, -1, 2, 0, 2, 1, 2, 2, 2, 3,
  765.                 3, -3, 3, -2, 3, -1, 3, 0, 3, 1, 3, 2, 3, 3,
  766.             ]
  767.         else:
  768.             raise NotImplementedError
  770.         offset = torch.Tensor(offset)
  771.         # offset[0::2] *= self.dilation[0]
  772.         # offset[1::2] *= self.dilation[1]
  773.         # a tuple of two ints – in which case, the first int is used for the height dimension, and the second int for the width dimension
  774.         self.register_buffer('dilated_offset', torch.Tensor(offset[None, None, ..., None, None]))  # B, G, 49, 1, 1
  775.         self.init_weights()
  777.         self.use_BFM = use_BFM
  778.         if use_BFM:
  779.             alpha = 8
  780.             BFM = np.zeros((self.in_channels, 1, self.kernel_size[0], self.kernel_size[0]))
  781.             for i in range(self.kernel_size[0]):
  782.                 for j in range(self.kernel_size[0]):
  783.                     point_1 = (i, j)
  784.                     point_2 = (self.kernel_size[0] // 2, self.kernel_size[0] // 2)
  785.                     dist = distance.euclidean(point_1, point_2)
  786.                     BFM[:, :, i, j] = alpha / (dist + alpha)
  787.             self.register_buffer('BFM', torch.Tensor(BFM))
  788.         if fs_cfg is not None:
  789.             if pre_fs:
  790.                 self.FS = FrequencySelection(self.in_channels - self.normal_conv_dim, **fs_cfg)
  791.             else:
  792.                 self.FS = FrequencySelection(1, **fs_cfg)  # use dilation
  793.         self.pre_fs = pre_fs
  795.     def freq_select(self, x):
  796.         if self.offset_freq is None:
  797.             pass
  798.         elif self.offset_freq in ('FLC_high', 'SLP_high'):
  799.             x - self.LP(x)
  800.         elif self.offset_freq in ('FLC_res', 'SLP_res'):
  801.             2 * x - self.LP(x)
  802.         else:
  803.             raise NotImplementedError
  804.         return x
  806.     def init_weights(self):
  807.         super().init_weights()
  808.         if hasattr(self, 'conv_offset'):
  809.             self.conv_offset.weight.data.zero_()
  810.             self.conv_offset.bias.data.fill_((self.dilation[0] - 1) / self.dilation[0] + 1e-4)
  811.             # self.conv_offset.bias.data.zero_()
  812.         # if hasattr(self, 'conv_offset_low'):
  813.         # self.conv_offset_low[1].weight.data.zero_()
  814.         if hasattr(self, 'conv_mask'):
  815.             self.conv_mask[1].weight.data.zero_()
  816.             self.conv_mask[1].bias.data.zero_()
  818.     def forward(self, x):
  819.         if self.normal_conv_dim > 0:
  820.             return self.mix_forward(x)
  821.         else:
  822.             return self.ad_forward(x)
  824.     def ad_forward(self, x):
  825.         if hasattr(self, 'FS') and self.pre_fs: x = self.FS(x)
  826.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  827.             c_att1, _, _, _, = self.OMNI_ATT1(x)
  828.             c_att2, _, _, _, = self.OMNI_ATT2(x)
  829.         elif hasattr(self, 'OMNI_ATT'):
  830.             c_att, _, _, _, = self.OMNI_ATT(x)
  831.         x = self.PAD(x)
  832.         offset = self.conv_offset(x)
  833.         offset = F.relu(offset, inplace=True) * self.dilation[0]  # ensure > 0
  834.         if hasattr(self, 'FS') and (self.pre_fs == False): x = self.FS(x, offset)
  835.         b, _, h, w = offset.shape
  836.         offset = offset.reshape(b, self.deform_groups, -1, h, w) * self.dilated_offset
  837.         offset = offset.reshape(b, -1, h, w)
  838.         mask = self.conv_mask(x)
  839.         mask = torch.sigmoid(mask)
  840.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  841.             offset = offset.reshape(1, -1, h, w)
  842.             # print(offset.max(), offset.min(), offset.mean())
  843.             mask = mask.reshape(1, -1, h, w)
  844.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  845.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, out, in, k, k
  846.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  847.             adaptive_weight = adaptive_weight_mean * (2 * c_att1.unsqueeze(2)) + (
  848.                         adaptive_weight - adaptive_weight_mean) * (2 * c_att2.unsqueeze(2))
  849.             adaptive_weight = adaptive_weight.reshape(-1, self.in_channels // self.groups, 3, 3)
  850.             if self.bias is not None:
  851.                 bias = self.bias.repeat(b)
  852.             else:
  853.                 bias = self.bias
  854.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight, bias,
  855.                                         self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  856.                                         # padding
  857.                                         (1, 1),  # dilation
  858.                                         self.groups * b, self.deform_groups * b)
  859.             return x.reshape(b, -1, h, w)
  860.         elif hasattr(self, 'OMNI_ATT'):
  861.             offset = offset.reshape(1, -1, h, w)
  862.             mask = mask.reshape(1, -1, h, w)
  863.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  864.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, out, in, k, k
  865.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  866.             if self.kernel_decompose == 'high':
  867.                 adaptive_weight = adaptive_weight_mean + (adaptive_weight - adaptive_weight_mean) * (
  868.                             2 * c_att.unsqueeze(2))
  869.             elif self.kernel_decompose == 'low':
  870.                 adaptive_weight = adaptive_weight_mean * (2 * c_att.unsqueeze(2)) + (
  871.                             adaptive_weight - adaptive_weight_mean)
  872.             adaptive_weight = adaptive_weight.reshape(-1, self.in_channels // self.groups, 3, 3)
  873.             if self.bias is not None:
  874.                 bias = self.bias.repeat(b)
  875.             else:
  876.                 bias = self.bias
  877.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight, bias,
  878.                                         self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  879.                                         # padding
  880.                                         (1, 1),  # dilation
  881.                                         self.groups * b, self.deform_groups * b)
  882.             return x.reshape(b, -1, h, w)
  883.         else:
  884.             return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias,
  885.                                            self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  886.                                            # padding
  887.                                            self.dilation, self.groups,
  888.                                            self.deform_groups)
  890.     def mix_forward(self, x):
  891.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  892.             c_att1, _, _, _, = self.OMNI_ATT1(x)
  893.             c_att2, _, _, _, = self.OMNI_ATT2(x)
  894.         elif hasattr(self, 'OMNI_ATT'):
  895.             c_att, _, _, _, = self.OMNI_ATT(x)
  896.         ori_x = x
  897.         normal_conv_x = ori_x[:, -self.normal_conv_dim:]  # ad:normal
  898.         x = ori_x[:, :-self.normal_conv_dim]
  899.         if hasattr(self, 'FS') and self.pre_fs: x = self.FS(x)
  900.         x = self.PAD(x)
  901.         offset = self.conv_offset(x)
  902.         if hasattr(self, 'FS') and (self.pre_fs == False): x = self.FS(x, F.interpolate(offset, x.shape[-2:],
  903.                                                                                         mode='bilinear', align_corners=(
  904.                                                                                                                                    x.shape[
  905.                                                                                                                                        -1] % 2) == 1))
  906.         # if hasattr(self, 'FS') and (self.pre_fs==False): x = self.FS(x, offset)
  907.         # offset = F.relu(offset, inplace=True) * self.dilation[0] # ensure > 0
  908.         offset[offset < 0] = offset[offset < 0].exp() - 1
  909.         b, _, h, w = offset.shape
  910.         offset = offset.reshape(b, self.deform_groups, -1, h, w) * self.dilated_offset
  911.         offset = offset.reshape(b, -1, h, w)
  912.         mask = self.conv_mask(x)
  913.         mask = torch.sigmoid(mask)
  914.         if hasattr(self, 'OMNI_ATT1') and hasattr(self, 'OMNI_ATT2'):
  915.             offset = offset.reshape(1, -1, h, w)
  916.             # print(offset.max(), offset.min(), offset.mean())
  917.             mask = mask.reshape(1, -1, h, w)
  918.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  919.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, out, in, k, k
  920.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  921.             adaptive_weight = adaptive_weight_mean * (2 * c_att1.unsqueeze(2)) + (
  922.                         adaptive_weight - adaptive_weight_mean) * (2 * c_att2.unsqueeze(2))
  923.             if self.bias is not None:
  924.                 bias = self.bias.repeat(b)
  925.             else:
  926.                 bias = self.bias
  927.             # adaptive_weight = adaptive_weight.reshape(-1, self.in_channels // self.groups, 3, 3)
  928.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight[:, :-self.normal_conv_dim].reshape(-1,
  929.                                                                                                             self.in_channels // self.groups,
  930.                                                                                                             self.kernel_size[
  931.                                                                                                                 0],
  932.                                                                                                             self.kernel_size[
  933.                                                                                                                 1]),
  934.                                         bias,
  935.                                         self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  936.                                         # padding
  937.                                         (1, 1),  # dilation
  938.                                         (self.in_channels - self.normal_conv_dim) * b, self.deform_groups * b)
  939.             x = x.reshape(b, -1, h, w)
  940.             normal_conv_x = F.conv2d(normal_conv_x.reshape(1, -1, h, w),
  941.                                      adaptive_weight[:, -self.normal_conv_dim:].reshape(-1,
  942.                                                                                         self.in_channels // self.groups,
  943.                                                                                         self.kernel_size[0],
  944.                                                                                         self.kernel_size[1]),
  945.                                      bias=bias, stride=self.stride, padding=self.padding, dilation=self.dilation,
  946.                                      groups=self.normal_conv_dim * b)
  947.             normal_conv_x = normal_conv_x.reshape(b, -1, h, w)
  948.             # return torch.cat([normal_conv_x, x], dim=1)
  949.             return torch.cat([x, normal_conv_x], dim=1)
  950.         elif hasattr(self, 'OMNI_ATT'):
  951.             offset = offset.reshape(1, -1, h, w)
  952.             mask = mask.reshape(1, -1, h, w)
  953.             x = x.reshape(1, -1, x.size(-2), x.size(-1))
  954.             adaptive_weight = self.weight.unsqueeze(0).repeat(b, 1, 1, 1, 1)  # b, out, in, k, k
  955.             adaptive_weight_mean = adaptive_weight.mean(dim=(-1, -2), keepdim=True)
  956.             if self.bias is not None:
  957.                 bias = self.bias.repeat(b)
  958.             else:
  959.                 bias = self.bias
  960.             if self.kernel_decompose == 'high':
  961.                 adaptive_weight = adaptive_weight_mean + (adaptive_weight - adaptive_weight_mean) * (
  962.                             2 * c_att.unsqueeze(2))
  963.             elif self.kernel_decompose == 'low':
  964.                 adaptive_weight = adaptive_weight_mean * (2 * c_att.unsqueeze(2)) + (
  965.                             adaptive_weight - adaptive_weight_mean)
  966.             x = modulated_deform_conv2d(x, offset, mask, adaptive_weight[:, :-self.normal_conv_dim].reshape(-1,
  967.                                                                                                             self.in_channels // self.groups,
  968.                                                                                                             self.kernel_size[
  969.                                                                                                                 0],
  970.                                                                                                             self.kernel_size[
  971.                                                                                                                 1]),
  972.                                         bias,
  973.                                         self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  974.                                         # padding
  975.                                         (1, 1),  # dilation
  976.                                         (self.in_channels - self.normal_conv_dim) * b, self.deform_groups * b)
  977.             x = x.reshape(b, -1, h, w)
  978.             normal_conv_x = F.conv2d(normal_conv_x.reshape(1, -1, h, w),
  979.                                      adaptive_weight[:, -self.normal_conv_dim:].reshape(-1,
  980.                                                                                         self.in_channels // self.groups,
  981.                                                                                         self.kernel_size[0],
  982.                                                                                         self.kernel_size[1]),
  983.                                      bias=bias, stride=self.stride, padding=self.padding, dilation=self.dilation,
  984.                                      groups=self.normal_conv_dim * b)
  985.             normal_conv_x = normal_conv_x.reshape(b, -1, h, w)
  986.             # return torch.cat([normal_conv_x, x], dim=1)
  987.             return torch.cat([x, normal_conv_x], dim=1)
  988.         else:
  989.             return modulated_deform_conv2d(x, offset, mask, self.weight, self.bias,
  990.                                            self.stride, self.padding if isinstance(self.PAD, nn.Identity) else 0,
  991.                                            # padding
  992.                                            self.dilation, self.groups,
  993.                                            self.deform_groups)
  996. def autopad(k, p=None, d=1):  # kernel, padding, dilation
  997.     """Pad to 'same' shape outputs."""
  998.     if d > 1:
  999.         k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
  1000.     if p is None:
  1001.         p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
  1002.     return p
  1005. class Conv(nn.Module):
  1006.     """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
  1008.     default_act = nn.SiLU()  # default activation
  1010.     def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
  1011.         """Initialize Conv layer with given arguments including activation."""
  1012.         super().__init__()
  1013.         self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
  1014.         self.bn = nn.BatchNorm2d(c2)
  1015.         self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
  1017.     def forward(self, x):
  1018.         """Apply convolution, batch normalization and activation to input tensor."""
  1019.         return self.act(self.bn(self.conv(x)))
  1021.     def forward_fuse(self, x):
  1022.         """Perform transposed convolution of 2D data."""
  1023.         return self.act(self.conv(x))
  1025. class DFL(nn.Module):
  1026.     """
  1027.     Integral module of Distribution Focal Loss (DFL).
  1029.     Proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
  1030.     """
  1032.     def __init__(self, c1=16):
  1033.         """Initialize a convolutional layer with a given number of input channels."""
  1034.         super().__init__()
  1035.         self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
  1036.         x = torch.arange(c1, dtype=torch.float)
  1037.         self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
  1038.         self.c1 = c1
  1040.     def forward(self, x):
  1041.         """Applies a transformer layer on input tensor 'x' and returns a tensor."""
  1042.         b, _, a = x.shape  # batch, channels, anchors
  1043.         return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
  1044.         # return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
  1046. class ADDWConvHead(nn.Module):
  1047.     """YOLOv8 Detect head for detection models."""
  1049.     dynamic = False  # force grid reconstruction
  1050.     export = False  # export mode
  1051.     end2end = False  # end2end
  1052.     max_det = 300  # max_det
  1053.     shape = None
  1054.     anchors = torch.empty(0)  # init
  1055.     strides = torch.empty(0)  # init
  1057.     def __init__(self, nc=80, ch=()):
  1058.         """Initializes the YOLOv8 detection layer with specified number of classes and channels."""
  1059.         super().__init__()
  1060.         self.nc = nc  # number of classes
  1061.         self.nl = len(ch)  # number of detection layers
  1062.         self.reg_max = 16  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
  1063.         self.no = nc + self.reg_max * 4  # number of outputs per anchor
  1064.         self.stride = torch.zeros(self.nl)  # strides computed during build
  1065.         c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))  # channels
  1066.         self.cv2 = nn.ModuleList(
  1067.             nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch
  1068.         )
  1069.         self.cv3 = nn.ModuleList(
  1070.             nn.Sequential(
  1071.                 nn.Sequential(AdaptiveDilatedDWConv(x, x, groups=math.gcd(x, x), kernel_size=3, stride=1, dilation=1), Conv(x, c3, 1)),
  1072.                 nn.Sequential(AdaptiveDilatedDWConv(c3, c3, groups=math.gcd(c3, c3), kernel_size=3, stride=1, dilation=1), Conv(c3, c3, 1)),
  1073.                 nn.Conv2d(c3, self.nc, 1),
  1074.             )
  1075.             for x in ch
  1076.         )
  1077.         self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
  1079.         if self.end2end:
  1080.             self.one2one_cv2 = copy.deepcopy(self.cv2)
  1081.             self.one2one_cv3 = copy.deepcopy(self.cv3)
  1083.     def forward(self, x):
  1084.         """Concatenates and returns predicted bounding boxes and class probabilities."""
  1085.         if self.end2end:
  1086.             return self.forward_end2end(x)
  1087.         for i in range(self.nl):
  1088.             x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
  1089.         if self.training:  # Training path
  1090.             return x
  1091.         y = self._inference(x)
  1092.         return y if self.export else (y, x)
  1094.     def forward_end2end(self, x):
  1095.         """
  1096.         Performs forward pass of the v10Detect module.
  1098.         Args:
  1099.             x (tensor): Input tensor.
  1101.         Returns:
  1102.             (dict, tensor): If not in training mode, returns a dictionary containing the outputs of both one2many and one2one detections.
  1103.                            If in training mode, returns a dictionary containing the outputs of one2many and one2one detections separately.
  1104.         """
  1105.         x_detach = [xi.detach() for xi in x]
  1106.         one2one = [
  1107.             torch.cat((self.one2one_cv2[i](x_detach[i]), self.one2one_cv3[i](x_detach[i])), 1) for i in range(self.nl)
  1108.         ]
  1109.         for i in range(self.nl):
  1110.             x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
  1111.         if self.training:  # Training path
  1112.             return {"one2many": x, "one2one": one2one}
  1114.         y = self._inference(one2one)
  1115.         y = self.postprocess(y.permute(0, 2, 1), self.max_det, self.nc)
  1116.         return y if self.export else (y, {"one2many": x, "one2one": one2one})
  1118.     def _inference(self, x):
  1119.         """Decode predicted bounding boxes and class probabilities based on multiple-level feature maps."""
  1120.         # Inference path
  1121.         shape = x[0].shape  # BCHW
  1122.         x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
  1123.         if self.dynamic or self.shape != shape:
  1124.             self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
  1125.             self.shape = shape
  1127.         if self.export and self.format in {"saved_model", "pb", "tflite", "edgetpu", "tfjs"}:  # avoid TF FlexSplitV ops
  1128.             box = x_cat[:, : self.reg_max * 4]
  1129.             cls = x_cat[:, self.reg_max * 4 :]
  1130.         else:
  1131.             box, cls = x_cat.split((self.reg_max * 4, self.nc), 1)
  1133.         if self.export and self.format in {"tflite", "edgetpu"}:
  1134.             # Precompute normalization factor to increase numerical stability
  1135.             # See https://github.com/ultralytics/ultralytics/issues/7371
  1136.             grid_h = shape[2]
  1137.             grid_w = shape[3]
  1138.             grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
  1139.             norm = self.strides / (self.stride[0] * grid_size)
  1140.             dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
  1141.         else:
  1142.             dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
  1144.         return torch.cat((dbox, cls.sigmoid()), 1)
  1146.     def bias_init(self):
  1147.         """Initialize Detect() biases, WARNING: requires stride availability."""
  1148.         m = self  # self.model[-1]  # Detect() module
  1149.         # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
  1150.         # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
  1151.         for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
  1152.             a[-1].bias.data[:] = 1.0  # box
  1153.             b[-1].bias.data[: m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
  1154.         if self.end2end:
  1155.             for a, b, s in zip(m.one2one_cv2, m.one2one_cv3, m.stride):  # from
  1156.                 a[-1].bias.data[:] = 1.0  # box
  1157.                 b[-1].bias.data[: m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
  1159.     def decode_bboxes(self, bboxes, anchors):
  1160.         """Decode bounding boxes."""
  1161.         return dist2bbox(bboxes, anchors, xywh=not self.end2end, dim=1)
  1163.     @staticmethod
  1164.     def postprocess(preds: torch.Tensor, max_det: int, nc: int = 80):
  1165.         """
  1166.         Post-processes YOLO model predictions.
  1168.         Args:
  1169.             preds (torch.Tensor): Raw predictions with shape (batch_size, num_anchors, 4 + nc) with last dimension
  1170.                 format [x, y, w, h, class_probs].
  1171.             max_det (int): Maximum detections per image.
  1172.             nc (int, optional): Number of classes. Default: 80.
  1174.         Returns:
  1175.             (torch.Tensor): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6) and last
  1176.                 dimension format [x, y, w, h, max_class_prob, class_index].
  1177.         """
  1178.         batch_size, anchors, _ = preds.shape  # i.e. shape(16,8400,84)
  1179.         boxes, scores = preds.split([4, nc], dim=-1)
  1180.         index = scores.amax(dim=-1).topk(min(max_det, anchors))[1].unsqueeze(-1)
  1181.         boxes = boxes.gather(dim=1, index=index.repeat(1, 1, 4))
  1182.         scores = scores.gather(dim=1, index=index.repeat(1, 1, nc))
  1183.         scores, index = scores.flatten(1).topk(min(max_det, anchors))
  1184.         i = torch.arange(batch_size)[..., None]  # batch indices
  1185.         return torch.cat([boxes[i, index // nc], scores[..., None], (index % nc)[..., None].float()], dim=-1)
  1189. if __name__ == "__main__":
  1190.     # Generating Sample image
  1191.     image1 = (1, 64, 160, 160)
  1192.     image2 = (1, 128, 80, 80)
  1193.     image3 = (1, 256, 40, 40)
  1195.     image1 = torch.rand(image1)
  1196.     image2 = torch.rand(image2)
  1197.     image3 = torch.rand(image3)
  1198.     image = [image1, image2, image3]
  1199.     channel = (64, 128, 256)
  1200.     # Model
  1201.     mobilenet_v1 = ADDWConvHead(nc=80,  ch=channel)
  1203.     out = mobilenet_v1(image)
  1204.     print(out)

四、添加教程

4.1 修改一

首先我们将上面的代码复制粘贴到'ultralytics/nn' 目录下新建一个py文件复制粘贴进去，具体名字自己来定.

4.2 修改二

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'( 用群内的文件的话已经有了无需新建) ，然后在其内部导入我们的检测头如下图所示。

4.3 修改三

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块( 用群内的文件的话已经有了无需重新导入直接开始第四步即可) ！

​​​

4.4 修改四

第四步我门找到如下文件'ultralytics/nn/tasks.py，找到如下的代码进行将检测头添加进去，这里给大家推荐个快速搜索的方法用ctrl+f然后搜索Detect然后就能快速查找了。

​​​​

4.5 修改五

同理

4.6 修改六

同理

4.7 修改七

这里有一些不一样，我们需要加一行代码 

  1.         else:
  2.             return 'detect'

为啥呢不一样，因为这里的m在代码执行过程中会将你的代码自动转换为小写，所以直接else方便一点，以后出现一些其它分割或者其它的教程的时候在提供其它的修改教程。

​​​​

4.8 修改八

同理.

​​​​

到此就修改完成了，大家可以复制下面的yaml文件运行。

五、FADDWConvHead检测头的yaml文件

此版本训练信息：YOLO11-FADDWConvHead summary: 446 layers, 2,667,632 parameters, 2,667,616 gradients, 6.6 GFLOPs 

  1. # Ultralytics YOLO 🚀, AGPL-3.0 license
  2. # YOLO11 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
  4. # Parameters
  5. nc: 80 # number of classes
  6. scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n'
  7.   # [depth, width, max_channels]
  8.   n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs
  9.   s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs
  10.   m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs
  11.   l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs
  12.   x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs
  14. # YOLO11n backbone
  15. backbone:
  16.   # [from, repeats, module, args]
  17.   - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  18.   - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  19.   - [-1, 2, C3k2, [256, False, 0.25]]
  20.   - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  21.   - [-1, 2, C3k2, [512, False, 0.25]]
  22.   - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  23.   - [-1, 2, C3k2, [512, True]]
  24.   - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  25.   - [-1, 2, C3k2, [1024, True]]
  26.   - [-1, 1, SPPF, [1024, 5]] # 9
  27.   - [-1, 2, C2PSA, [1024]] # 10
  29. # YOLO11n head
  30. head:
  31.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  32.   - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  33.   - [-1, 2, C3k2, [512, False]] # 13
  35.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  36.   - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  37.   - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)
  39.   - [-1, 1, Conv, [256, 3, 2]]
  40.   - [[-1, 13], 1, Concat, [1]] # cat head P4
  41.   - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)
  43.   - [-1, 1, Conv, [512, 3, 2]]
  44.   - [[-1, 10], 1, Concat, [1]] # cat head P5
  45.   - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)
  47.   - [[16, 19, 22], 1, FADDWConvHead, [nc]] # Detect(P3, P4, P5)

六、完美运行记录

最后提供一下完美运行的图片。

​​

​​

七、本文总结

到此本文的正式分享内容就结束了，在这里给大家推荐我的YOLOv11改进有效涨点专栏，本专栏目前为新开的平均质量分98分，后期我会根据各种最新的前沿顶会进行论文复现，也会对一些老的改进机制进行补充，如果大家觉得本文帮助到你了，订阅本专栏，关注后续更多的更新~