• YOLOv8改进 | 如何在网络结构中添加注意力机制、C2f、卷积、Neck、检测头

    一、本文介绍

    本篇文章的内容是在大家得到一个改进版本的C2f一个新的注意力机制、或者一个新的卷积模块、或者是检测头的时候如何替换我们YOLOv8模型中的原有的模块，从而用你的模块去进行训练模型或者检测。因为最近开了一个专栏里面涉及到挺多改进的地方，不能每篇文章都去讲解一遍如何修改，就想着在这里单独出一期文章进行一个总结性教程，大家可以从我的其它文章中拿到修改后的代码，从这篇文章学会如何去添加到你的模型结构中去。

    YOLOv8专栏：YOLOv8改进有效涨点专栏->持续复现各种最新机制

    本文的讲解举例都以最新的YOLOv8的目录结构为例，老版本的其实方法都一样只是目录构造不一样找到同样的文件名即可。 

    适用对象->本文适合那些拿到源码却不知道如何添加到网络结构中的朋友

    目录

    一、本文介绍

    二、导入修改内容

    2.1创建新文件导入新模块

    2.1.1情况一

    2.1.2情况二 

    三、Conv模块

    3.1修改一

    3.2修改二

    3.3修改三 

    四、C2f、Bottleneck模块

    4.1修改一

    4.2步骤二 

    4.3修改三

    4.4修改四

    4.5修改五

    4.6修改六

    4.7修改七

    四、注意力机制

    4.1修改一 

    4.2修改二

    4.3修改三

    4.2.1有参数的注意力机制修改

    4.2.2无参数的注意力机制修改

    4.4配置注意力机制

    五、Neck部分

    六、检测头

    七、损失函数

    二、导入修改内容

    大家拿到任何一个代码，想要加入到模型的内部，我们都需要先将其导入到模型的内部，才可以将其添加到模型的结构中去，下面的代码是一个ODConv，和我创建的一个ODConv_yolo的类(官方的代码报错进行一定的处理想知道为啥可以看我单独讲解它的博客)， 我们先拿其进行举例。

    1. import torch
    2. import torch.nn as nn
    3. import torch.nn.functional as F
    4. import torch.autograd
    5.  
    6.  
    7. class Attention(nn.Module):
    8.     def __init__(self, in_planes, out_planes, kernel_size, groups=1, reduction=0.0625, kernel_num=4, min_channel=16):
    9.         super(Attention, self).__init__()
    10.         attention_channel = max(int(in_planes * reduction), min_channel)
    11.         self.kernel_size = kernel_size
    12.         self.kernel_num = kernel_num
    13.         self.temperature = 1.0
    14.  
    15.         self.avgpool = nn.AdaptiveAvgPool2d(1)
    16.         self.fc = nn.Conv2d(in_planes, attention_channel, 1, bias=False)
    17.         self.bn = nn.BatchNorm2d(attention_channel)
    18.         self.relu = nn.ReLU(inplace=True)
    19.  
    20.         self.channel_fc = nn.Conv2d(attention_channel, in_planes, 1, bias=True)
    21.         self.func_channel = self.get_channel_attention
    22.  
    23.         if in_planes == groups and in_planes == out_planes:  # depth-wise convolution
    24.             self.func_filter = self.skip
    25.         else:
    26.             self.filter_fc = nn.Conv2d(attention_channel, out_planes, 1, bias=True)
    27.             self.func_filter = self.get_filter_attention
    28.  
    29.         if kernel_size == 1:  # point-wise convolution
    30.             self.func_spatial = self.skip
    31.         else:
    32.             self.spatial_fc = nn.Conv2d(attention_channel, kernel_size * kernel_size, 1, bias=True)
    33.             self.func_spatial = self.get_spatial_attention
    34.  
    35.         if kernel_num == 1:
    36.             self.func_kernel = self.skip
    37.         else:
    38.             self.kernel_fc = nn.Conv2d(attention_channel, kernel_num, 1, bias=True)
    39.             self.func_kernel = self.get_kernel_attention
    40.  
    41.         self._initialize_weights()
    42.  
    43.     def _initialize_weights(self):
    44.         for m in self.modules():
    45.             if isinstance(m, nn.Conv2d):
    46.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
    47.                 if m.bias is not None:
    48.                     nn.init.constant_(m.bias, 0)
    49.             if isinstance(m, nn.BatchNorm2d):
    50.                 nn.init.constant_(m.weight, 1)
    51.                 nn.init.constant_(m.bias, 0)
    52.  
    53.     def update_temperature(self, temperature):
    54.         self.temperature = temperature
    55.  
    56.     @staticmethod
    57.     def skip(_):
    58.         return 1.0
    59.  
    60.     def get_channel_attention(self, x):
    61.         channel_attention = torch.sigmoid(self.channel_fc(x).view(x.size(0), -1, 1, 1) / self.temperature)
    62.         return channel_attention
    63.  
    64.     def get_filter_attention(self, x):
    65.         filter_attention = torch.sigmoid(self.filter_fc(x).view(x.size(0), -1, 1, 1) / self.temperature)
    66.         return filter_attention
    67.  
    68.     def get_spatial_attention(self, x):
    69.         spatial_attention = self.spatial_fc(x).view(x.size(0), 1, 1, 1, self.kernel_size, self.kernel_size)
    70.         spatial_attention = torch.sigmoid(spatial_attention / self.temperature)
    71.         return spatial_attention
    72.  
    73.     def get_kernel_attention(self, x):
    74.         kernel_attention = self.kernel_fc(x).view(x.size(0), -1, 1, 1, 1, 1)
    75.         kernel_attention = F.softmax(kernel_attention / self.temperature, dim=1)
    76.         return kernel_attention
    77.  
    78.     def forward(self, x):
    79.         x = self.avgpool(x)
    80.         x = self.fc(x)
    81.         # x = self.bn(x) # 在外面我提供了一个bn这里会报错
    82.         x = self.relu(x)
    83.         return self.func_channel(x), self.func_filter(x), self.func_spatial(x), self.func_kernel(x)
    84.  
    85.  
    86. class ODConv2d(nn.Module):
    87.     def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=1, dilation=1, groups=1,
    88.                  reduction=0.0625, kernel_num=4):
    89.         super(ODConv2d, self).__init__()
    90.         kernel_size = kernel_size[0]
    91.         in_planes = in_planes
    92.         self.in_planes = in_planes
    93.         self.out_planes = out_planes
    94.         self.kernel_size = kernel_size
    95.         self.stride = stride
    96.         self.padding = padding
    97.         self.dilation = dilation
    98.         self.groups = groups
    99.         self.kernel_num = kernel_num
    100.         self.attention = Attention(in_planes, out_planes, kernel_size, groups=groups,
    101.                                    reduction=reduction, kernel_num=kernel_num)
    102.         self.weight = nn.Parameter(torch.randn(kernel_num, out_planes, in_planes//groups, kernel_size, kernel_size),
    103.                                    requires_grad=True)
    104.         self._initialize_weights()
    105.  
    106.         if self.kernel_size == 1 and self.kernel_num == 1:
    107.             self._forward_impl = self._forward_impl_pw1x
    108.         else:
    109.             self._forward_impl = self._forward_impl_common
    110.  
    111.     def _initialize_weights(self):
    112.         for i in range(self.kernel_num):
    113.             nn.init.kaiming_normal_(self.weight[i], mode='fan_out', nonlinearity='relu')
    114.  
    115.     def update_temperature(self, temperature):
    116.         self.attention.update_temperature(temperature)
    117.  
    118.     def _forward_impl_common(self, x):
    119.         # Multiplying channel attention (or filter attention) to weights and feature maps are equivalent,
    120.         # while we observe that when using the latter method the models will run faster with less gpu memory cost.
    121.         channel_attention, filter_attention, spatial_attention, kernel_attention = self.attention(x)
    122.         batch_size, in_planes, height, width = x.size()
    123.         x = x * channel_attention
    124.         x = x.reshape(1, -1, height, width)
    125.         aggregate_weight = spatial_attention * kernel_attention * self.weight.unsqueeze(dim=0)
    126.         aggregate_weight = torch.sum(aggregate_weight, dim=1).view(
    127.             [-1, self.in_planes // self.groups, self.kernel_size, self.kernel_size])
    128.         output = F.conv2d(x, weight=aggregate_weight, bias=None, stride=self.stride, padding=self.padding,
    129.                           dilation=self.dilation, groups=self.groups * batch_size)
    130.         output = output.view(batch_size, self.out_planes, output.size(-2), output.size(-1))
    131.         output = output * filter_attention
    132.         return output
    133.  
    134.     def _forward_impl_pw1x(self, x):
    135.         channel_attention, filter_attention, spatial_attention, kernel_attention = self.attention(x)
    136.         x = x * channel_attention
    137.         output = F.conv2d(x, weight=self.weight.squeeze(dim=0), bias=None, stride=self.stride, padding=self.padding,
    138.                           dilation=self.dilation, groups=self.groups)
    139.         output = output * filter_attention
    140.         return output
    141.  
    142.     def forward(self, x):
    143.         return self._forward_impl(x)
    144.

    拿到这种代码之后，一般都很长，有一些博主推荐直接将其复制粘贴到YOLOv8的"ultralytics/nn/modules/conv.py"或者"ultralytics/nn/modules/block.py"目录下面，这种方法可不可以？答案是可以的，但是我建议大家最好新建一个文件在conv.py的同级目录下，为什么这么做，因为我们修改的模块越来越多，你往conv.py文件或则block.py文件里面加的代码越来越多很容易就把代码改崩溃了，最后就跌卸载进行重新下载包，我们通过建立文件导入其中类的形式，如果我们不用了，也不会对我们的代码做出任何影响，实在不行把新建立的文件删除了都可以，下面开始进行实际操作的讲解。

    2.1创建新文件导入新模块

    我们将我们得到的任何一个Conv或者想要修改的任何一个模块都可以像下面的图片一样直接建立一个文件复制粘贴进去即可。

    建立好上面的文件之后，我们此时呢有两种情况，一周呢官方的代码可以直接使用，另一种呢需要进行一定的处理，我们下面分别进行讲解两种情况。

    2.1.1情况一

    这种情况是官方的代码可以直接使用，此时我们直接修改"ultralytics/nn/modules/__init__.py"文件就可以了，修改如下->

    2.1.2情况二 

    另一种情况(绝大多数)：官方的代码不能直接使用我们本文的例子ODConv就是这种情况，所以我们需要对其进行一定的处理，我们找到如下的文件->"ultralytics/nn/modules/conv.py"对其进行修改如下->  

    修改一、导入模块

    修改二、将额外处理代码添加至conv模块 

    将如下代码添加至该文件中的末尾处-> 

    1.  
    2. class ODConv2d_yolo(nn.Module):
    3.     def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, groups=1, dilation=1):
    4.         super().__init__()
    5.         self.conv = Conv(in_channels, out_channels, k=1)
    6.         self.dcnv3 = ODConv2d(out_channels,out_channels, kernel_size=kernel_size, stride=stride, groups=groups,
    7.                                    dilation=dilation)
    8.         self.bn = nn.BatchNorm2d(out_channels)
    9.         self.gelu = nn.GELU()
    10.  
    11.     def forward(self, x):
    12.         x = self.conv(x)
    13.  
    14.         x = self.dcnv3(x)
    15.  
    16.         x = self.gelu(self.bn(x))
    17.         return x

    修改三、配置头文件

    修改如下->

    修改四 、重复情况一的步骤

    修改"ultralytics/nn/modules/__init__.py"文件如下

    总结：通过建立文件这种方法导入想要加入到模型中的模块(这里举例的是ODConv2d)其已经在我们新创建的.py文件中定义好了然后直接导过来就可以用了，从而不修改原有的conv.py文件就做到了，这样就算我们随时不用了，直接删除文件然后需要改的地方也很直观，否则时间久了代码早晚跌崩溃。

    三、Conv模块

    上面我们已经把定义好的卷积模块代码中了，此时我们还需要配置其位置，当然不同的模块导入的方式也有可能略有不同。

    3.1修改一

    我们找到如下的文件"ultralytics/nn/tasks.py"，图片如下->

    我们先把我们在上面"ultralytics/nn/modules/__init__.py" 文件的函数头中导入的类，在下面的地方导入进"ultralytics/nn/tasks.py"文件中，修改内容如下-> 

    3.2修改二

    我们在这个文件中找到一个方法(def定义的就叫方法)，因为其代码很长，我们一行一行搜索很麻烦，我们适用文件搜索功能(快捷键Ctrl + F)，弹出快捷栏如下->

    我们搜索下面这个代码"parse_model" 然后进行翻滚很容易就找到了下面的部分，同时进行红框内部的修改

    3.3修改三 

    到此我们就已经将我们定义的三个模块添加到我们的模型中了，已经可以修改yaml文件进行网络结构的配置了，我们找到该文件"ultralytics/cfg/models/v8/yolov8.yaml"进行配置。

    我们可以在其中的任何一个位置进行替换，当然我们的替换要符合逻辑，类似于之前这个位置是Conv那么你可以将你修改的卷积替换上，之前这个位置是C2f那么你就将修改后的C2f替换上。

    我们在yaml文件中进行了如下修改。

    到此我们就配置完成了此时进行训练就可以开始训练了~

    四、C2f、Bottleneck模块

    下面我们拿修改后的C2f、和Bottleneck举例，这两个模块定义在该文件中"ultralytics/nn/modules/block.py"，所以如果我们想添加修改后的C2f和Bottleneck(这俩一般配套使用)，就需要在该文件中进行修改，修改步骤如下->

    4.1修改一

    找到该文件"ultralytics/nn/modules/block.py"，进行如下修改->

    4.2步骤二 

    添加修改后的C2f和Bottleneck模块，这里起名为C2f_ODConv和Bottleneck_ODConv，

    1. class Bottleneck_ODConv(nn.Module):
    2.     """Standard bottleneck."""
    3.  
    4.     def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
    5.         """Initializes a bottleneck module with given input/output channels, shortcut option, group, kernels, and
    6.         expansion.
    7.         """
    8.         super().__init__()
    9.         c_ = int(c2 * e)  # hidden channels
    10.         self.cv1 = Conv(c1, c_, k[0], 1)
    11.         self.cv2 = ODConv2d_yolo(c_, c2, k[1], 1, groups=g)
    12.         self.add = shortcut and c1 == c2
    13.  
    14.     def forward(self, x):
    15.         """'forward()' applies the YOLO FPN to input data."""
    16.         return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
    17.  
    18.  
    19. class C2f_ODConv(nn.Module):
    20.     """Faster Implementation of CSP Bottleneck with 2 convolutions."""
    21.  
    22.     def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
    23.         """Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups,
    24.         expansion.
    25.         """
    26.         super().__init__()
    27.         self.c = int(c2 * e)  # hidden channels
    28.         self.cv1 = Conv(c1, 2 * self.c, 1, 1)
    29.         self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)
    30.         self.m = nn.ModuleList(Bottleneck_ODConv(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
    31.  
    32.     def forward(self, x):
    33.         """Forward pass through C2f layer."""
    34.         y = list(self.cv1(x).chunk(2, 1))
    35.         y.extend(m(y[-1]) for m in self.m)
    36.         return self.cv2(torch.cat(y, 1))
    37.  
    38.     def forward_split(self, x):
    39.         """Forward pass using split() instead of chunk()."""
    40.         y = list(self.cv1(x).split((self.c, self.c), 1))
    41.         y.extend(m(y[-1]) for m in self.m)
    42.         return self.cv2(torch.cat(y, 1))

    将以上代码复制到文件"ultralytics/nn/modules/block.py"的末尾， 

    4.3修改三

    修改头文件如下->

    4.4修改四

    找到文件"ultralytics/nn/modules/__init__.py"，修改如下->

    4.5修改五

    找到该文件我们找到如下的文件"ultralytics/nn/tasks.py"进行修改(其实和卷积模块的一模一样)，

    4.6修改六

    我们在这个文件中找到一个方法(def定义的就叫方法)，因为其代码很长，我们一行一行搜索很麻烦，我们适用文件搜索功能(快捷键Ctrl + F)，弹出快捷栏如下->

    我们搜索下面这个代码"parse_model" 然后进行翻滚很容易就找到了下面的部分，同时进行红框内部的修改

    4.7修改七

    到此我们就已经将我们定义的三个模块添加到我们的模型中了，已经可以修改yaml文件进行网络结构的配置了，我们找到该文件"ultralytics/cfg/models/v8/yolov8.yaml"进行配置。

    我们可以在其中的任何一个位置进行替换，当然我们的替换要符合逻辑，类似于之前这个位置是Conv那么你可以将你修改的卷积替换上，之前这个位置是C2f那么你就将修改后的C2f替换上。

    在yaml文件中进行了如下修改。

    到此就完成了修改C2f和Bottleneck模块了，已经可以开始进行训练了~

    至于修改这个ODConv的 效果如何可以看我的其它博客里面有详细的讲解~

    四、注意力机制

    修改注意力机制的部分其实和上面都是类似只是在修改如下文件的时候有点不一样"ultralytics/nn/tasks.py"，但是需要注意的是注意力机制分为两种，一种是有参数的注意力机制我们需要像其中传入参数，一种是无参数的注意力机制这两种机制的添加呢稍微有一些不同，我会在下面进行标注大家仔细看

    4.1修改一 

    这里我们拿Biformer注意力机制为例(我们拿有参数的注意力机制为例)，首先我们找到该目录'ultralytics/nn/modules'该目录的构造如下->

    我们在其中创建一个名字为Biformer的py文件如图所示，我们在其中复制如下代码即可

    1. """
    2. Bi-Level Routing Attention.
    3. """
    4. from typing import Tuple, Optional
    5. import torch
    6. import torch.nn as nn
    7. import torch.nn.functional as F
    8. from einops import rearrange
    9. from torch import Tensor, LongTensor
    10.  
    11. class TopkRouting(nn.Module):
    12.     """
    13.     differentiable topk routing with scaling
    14.     Args:
    15.         qk_dim: int, feature dimension of query and key
    16.         topk: int, the 'topk'
    17.         qk_scale: int or None, temperature (multiply) of softmax activation
    18.         with_param: bool, wether inorporate learnable params in routing unit
    19.         diff_routing: bool, wether make routing differentiable
    20.         soft_routing: bool, wether make output value multiplied by routing weights
    21.     """
    22.  
    23.     def __init__(self, qk_dim, topk=4, qk_scale=None, param_routing=False, diff_routing=False):
    24.         super().__init__()
    25.         self.topk = topk
    26.         self.qk_dim = qk_dim
    27.         self.scale = qk_scale or qk_dim ** -0.5
    28.         self.diff_routing = diff_routing
    29.         # TODO: norm layer before/after linear?
    30.         self.emb = nn.Linear(qk_dim, qk_dim) if param_routing else nn.Identity()
    31.         # routing activation
    32.         self.routing_act = nn.Softmax(dim=-1)
    33.  
    34.     def forward(self, query: Tensor, key: Tensor) -> Tuple[Tensor]:
    35.         """
    36.         Args:
    37.             q, k: (n, p^2, c) tensor
    38.         Return:
    39.             r_weight, topk_index: (n, p^2, topk) tensor
    40.         """
    41.         if not self.diff_routing:
    42.             query, key = query.detach(), key.detach()
    43.         query_hat, key_hat = self.emb(query), self.emb(key)  # per-window pooling -> (n, p^2, c)
    44.         attn_logit = (query_hat * self.scale) @ key_hat.transpose(-2, -1)  # (n, p^2, p^2)
    45.         topk_attn_logit, topk_index = torch.topk(attn_logit, k=self.topk, dim=-1)  # (n, p^2, k), (n, p^2, k)
    46.         r_weight = self.routing_act(topk_attn_logit)  # (n, p^2, k)
    47.  
    48.         return r_weight, topk_index
    49.  
    50.  
    51. class KVGather(nn.Module):
    52.     def __init__(self, mul_weight='none'):
    53.         super().__init__()
    54.         assert mul_weight in ['none', 'soft', 'hard']
    55.         self.mul_weight = mul_weight
    56.  
    57.     def forward(self, r_idx: Tensor, r_weight: Tensor, kv: Tensor):
    58.         """
    59.         r_idx: (n, p^2, topk) tensor
    60.         r_weight: (n, p^2, topk) tensor
    61.         kv: (n, p^2, w^2, c_kq+c_v)
    63.         Return:
    64.             (n, p^2, topk, w^2, c_kq+c_v) tensor
    65.         """
    66.         # select kv according to routing index
    67.         n, p2, w2, c_kv = kv.size()
    68.         topk = r_idx.size(-1)
    69.         # print(r_idx.size(), r_weight.size())
    70.         # FIXME: gather consumes much memory (topk times redundancy), write cuda kernel?
    71.         topk_kv = torch.gather(kv.view(n, 1, p2, w2, c_kv).expand(-1, p2, -1, -1, -1),
    72.                                # (n, p^2, p^2, w^2, c_kv) without mem cpy
    73.                                dim=2,
    74.                                index=r_idx.view(n, p2, topk, 1, 1).expand(-1, -1, -1, w2, c_kv)
    75.                                # (n, p^2, k, w^2, c_kv)
    76.                                )
    77.  
    78.         if self.mul_weight == 'soft':
    79.             topk_kv = r_weight.view(n, p2, topk, 1, 1) * topk_kv  # (n, p^2, k, w^2, c_kv)
    80.         elif self.mul_weight == 'hard':
    81.             raise NotImplementedError('differentiable hard routing TBA')
    82.         # else: #'none'
    83.         #     topk_kv = topk_kv # do nothing
    84.  
    85.         return topk_kv
    86.  
    87.  
    88. class QKVLinear(nn.Module):
    89.     def __init__(self, dim, qk_dim, bias=True):
    90.         super().__init__()
    91.         self.dim = dim
    92.         self.qk_dim = qk_dim
    93.         self.qkv = nn.Linear(dim, qk_dim + qk_dim + dim, bias=bias)
    94.  
    95.     def forward(self, x):
    96.         q, kv = self.qkv(x).split([self.qk_dim, self.qk_dim + self.dim], dim=-1)
    97.         return q, kv
    98.         # q, k, v = self.qkv(x).split([self.qk_dim, self.qk_dim, self.dim], dim=-1)
    99.         # return q, k, v
    100.  
    101.  
    102. class BiLevelRoutingAttention(nn.Module):
    103.     """
    104.     n_win: number of windows in one side (so the actual number of windows is n_win*n_win)
    105.     kv_per_win: for kv_downsample_mode='ada_xxxpool' only, number of key/values per window. Similar to n_win, the actual number is kv_per_win*kv_per_win.
    106.     topk: topk for window filtering
    107.     param_attention: 'qkvo'-linear for q,k,v and o, 'none': param free attention
    108.     param_routing: extra linear for routing
    109.     diff_routing: wether to set routing differentiable
    110.     soft_routing: wether to multiply soft routing weights
    111.     """
    112.  
    113.     def __init__(self, dim, n_win=7, num_heads=8, qk_dim=None, qk_scale=None,
    114.                  kv_per_win=4, kv_downsample_ratio=4, kv_downsample_kernel=None, kv_downsample_mode='identity',
    115.                  topk=4, param_attention="qkvo", param_routing=False, diff_routing=False, soft_routing=False,
    116.                  side_dwconv=3,
    117.                  auto_pad=True):
    118.         super().__init__()
    119.         # local attention setting
    120.         self.dim = dim
    121.         self.n_win = n_win  # Wh, Ww
    122.         self.num_heads = num_heads
    123.         self.qk_dim = qk_dim or dim
    124.         assert self.qk_dim % num_heads == 0 and self.dim % num_heads == 0, 'qk_dim and dim must be divisible by num_heads!'
    125.         self.scale = qk_scale or self.qk_dim ** -0.5
    126.  
    127.         ################side_dwconv (i.e. LCE in ShuntedTransformer)###########
    128.         self.lepe = nn.Conv2d(dim, dim, kernel_size=side_dwconv, stride=1, padding=side_dwconv // 2,
    129.                               groups=dim) if side_dwconv > 0 else \
    130.             lambda x: torch.zeros_like(x)
    131.  
    132.         ################ global routing setting #################
    133.         self.topk = topk
    134.         self.param_routing = param_routing
    135.         self.diff_routing = diff_routing
    136.         self.soft_routing = soft_routing
    137.         # router
    138.         assert not (self.param_routing and not self.diff_routing)  # cannot be with_param=True and diff_routing=False
    139.         self.router = TopkRouting(qk_dim=self.qk_dim,
    140.                                   qk_scale=self.scale,
    141.                                   topk=self.topk,
    142.                                   diff_routing=self.diff_routing,
    143.                                   param_routing=self.param_routing)
    144.         if self.soft_routing:  # soft routing, always diffrentiable (if no detach)
    145.             mul_weight = 'soft'
    146.         elif self.diff_routing:  # hard differentiable routing
    147.             mul_weight = 'hard'
    148.         else:  # hard non-differentiable routing
    149.             mul_weight = 'none'
    150.         self.kv_gather = KVGather(mul_weight=mul_weight)
    151.  
    152.         # qkv mapping (shared by both global routing and local attention)
    153.         self.param_attention = param_attention
    154.         if self.param_attention == 'qkvo':
    155.             self.qkv = QKVLinear(self.dim, self.qk_dim)
    156.             self.wo = nn.Linear(dim, dim)
    157.         elif self.param_attention == 'qkv':
    158.             self.qkv = QKVLinear(self.dim, self.qk_dim)
    159.             self.wo = nn.Identity()
    160.         else:
    161.             raise ValueError(f'param_attention mode {self.param_attention} is not surpported!')
    162.  
    163.         self.kv_downsample_mode = kv_downsample_mode
    164.         self.kv_per_win = kv_per_win
    165.         self.kv_downsample_ratio = kv_downsample_ratio
    166.         self.kv_downsample_kenel = kv_downsample_kernel
    167.         if self.kv_downsample_mode == 'ada_avgpool':
    168.             assert self.kv_per_win is not None
    169.             self.kv_down = nn.AdaptiveAvgPool2d(self.kv_per_win)
    170.         elif self.kv_downsample_mode == 'ada_maxpool':
    171.             assert self.kv_per_win is not None
    172.             self.kv_down = nn.AdaptiveMaxPool2d(self.kv_per_win)
    173.         elif self.kv_downsample_mode == 'maxpool':
    174.             assert self.kv_downsample_ratio is not None
    175.             self.kv_down = nn.MaxPool2d(self.kv_downsample_ratio) if self.kv_downsample_ratio > 1 else nn.Identity()
    176.         elif self.kv_downsample_mode == 'avgpool':
    177.             assert self.kv_downsample_ratio is not None
    178.             self.kv_down = nn.AvgPool2d(self.kv_downsample_ratio) if self.kv_downsample_ratio > 1 else nn.Identity()
    179.         elif self.kv_downsample_mode == 'identity':  # no kv downsampling
    180.             self.kv_down = nn.Identity()
    181.         elif self.kv_downsample_mode == 'fracpool':
    182.             # assert self.kv_downsample_ratio is not None
    183.             # assert self.kv_downsample_kenel is not None
    184.             # TODO: fracpool
    185.             # 1. kernel size should be input size dependent
    186.             # 2. there is a random factor, need to avoid independent sampling for k and v
    187.             raise NotImplementedError('fracpool policy is not implemented yet!')
    188.         elif kv_downsample_mode == 'conv':
    189.             # TODO: need to consider the case where k != v so that need two downsample modules
    190.             raise NotImplementedError('conv policy is not implemented yet!')
    191.         else:
    192.             raise ValueError(f'kv_down_sample_mode {self.kv_downsaple_mode} is not surpported!')
    193.  
    194.         # softmax for local attention
    195.         self.attn_act = nn.Softmax(dim=-1)
    196.  
    197.         self.auto_pad = auto_pad
    198.  
    199.     def forward(self, x, ret_attn_mask=False):
    200.         """
    201.         x: NHWC tensor
    203.         Return:
    204.             NHWC tensor
    205.         """
    206.         x = rearrange(x, "n c h w -> n h w c")
    207.         # NOTE: use padding for semantic segmentation
    208.         ###################################################
    209.         if self.auto_pad:
    210.             N, H_in, W_in, C = x.size()
    211.  
    212.             pad_l = pad_t = 0
    213.             pad_r = (self.n_win - W_in % self.n_win) % self.n_win
    214.             pad_b = (self.n_win - H_in % self.n_win) % self.n_win
    215.             x = F.pad(x, (0, 0,  # dim=-1
    216.                           pad_l, pad_r,  # dim=-2
    217.                           pad_t, pad_b))  # dim=-3
    218.             _, H, W, _ = x.size()  # padded size
    219.         else:
    220.             N, H, W, C = x.size()
    221.             assert H % self.n_win == 0 and W % self.n_win == 0  #
    222.         ###################################################
    223.  
    224.         # patchify, (n, p^2, w, w, c), keep 2d window as we need 2d pooling to reduce kv size
    225.         x = rearrange(x, "n (j h) (i w) c -> n (j i) h w c", j=self.n_win, i=self.n_win)
    226.  
    227.         #################qkv projection###################
    228.         # q: (n, p^2, w, w, c_qk)
    229.         # kv: (n, p^2, w, w, c_qk+c_v)
    230.         # NOTE: separte kv if there were memory leak issue caused by gather
    231.         q, kv = self.qkv(x)
    232.  
    233.         # pixel-wise qkv
    234.         # q_pix: (n, p^2, w^2, c_qk)
    235.         # kv_pix: (n, p^2, h_kv*w_kv, c_qk+c_v)
    236.         q_pix = rearrange(q, 'n p2 h w c -> n p2 (h w) c')
    237.         kv_pix = self.kv_down(rearrange(kv, 'n p2 h w c -> (n p2) c h w'))
    238.         kv_pix = rearrange(kv_pix, '(n j i) c h w -> n (j i) (h w) c', j=self.n_win, i=self.n_win)
    239.  
    240.         q_win, k_win = q.mean([2, 3]), kv[..., 0:self.qk_dim].mean(
    241.             [2, 3])  # window-wise qk, (n, p^2, c_qk), (n, p^2, c_qk)
    242.  
    243.         ##################side_dwconv(lepe)##################
    244.         # NOTE: call contiguous to avoid gradient warning when using ddp
    245.         lepe = self.lepe(rearrange(kv[..., self.qk_dim:], 'n (j i) h w c -> n c (j h) (i w)', j=self.n_win,
    246.                                    i=self.n_win).contiguous())
    247.         lepe = rearrange(lepe, 'n c (j h) (i w) -> n (j h) (i w) c', j=self.n_win, i=self.n_win)
    248.  
    249.         ############ gather q dependent k/v #################
    250.  
    251.         r_weight, r_idx = self.router(q_win, k_win)  # both are (n, p^2, topk) tensors
    252.  
    253.         kv_pix_sel = self.kv_gather(r_idx=r_idx, r_weight=r_weight, kv=kv_pix)  # (n, p^2, topk, h_kv*w_kv, c_qk+c_v)
    254.         k_pix_sel, v_pix_sel = kv_pix_sel.split([self.qk_dim, self.dim], dim=-1)
    255.         # kv_pix_sel: (n, p^2, topk, h_kv*w_kv, c_qk)
    256.         # v_pix_sel: (n, p^2, topk, h_kv*w_kv, c_v)
    257.  
    258.         ######### do attention as normal ####################
    259.         k_pix_sel = rearrange(k_pix_sel, 'n p2 k w2 (m c) -> (n p2) m c (k w2)',
    260.                               m=self.num_heads)  # flatten to BMLC, (n*p^2, m, topk*h_kv*w_kv, c_kq//m) transpose here?
    261.         v_pix_sel = rearrange(v_pix_sel, 'n p2 k w2 (m c) -> (n p2) m (k w2) c',
    262.                               m=self.num_heads)  # flatten to BMLC, (n*p^2, m, topk*h_kv*w_kv, c_v//m)
    263.         q_pix = rearrange(q_pix, 'n p2 w2 (m c) -> (n p2) m w2 c',
    264.                           m=self.num_heads)  # to BMLC tensor (n*p^2, m, w^2, c_qk//m)
    265.  
    266.         # param-free multihead attention
    267.         attn_weight = (
    268.                                   q_pix * self.scale) @ k_pix_sel  # (n*p^2, m, w^2, c) @ (n*p^2, m, c, topk*h_kv*w_kv) -> (n*p^2, m, w^2, topk*h_kv*w_kv)
    269.         attn_weight = self.attn_act(attn_weight)
    270.         out = attn_weight @ v_pix_sel  # (n*p^2, m, w^2, topk*h_kv*w_kv) @ (n*p^2, m, topk*h_kv*w_kv, c) -> (n*p^2, m, w^2, c)
    271.         out = rearrange(out, '(n j i) m (h w) c -> n (j h) (i w) (m c)', j=self.n_win, i=self.n_win,
    272.                         h=H // self.n_win, w=W // self.n_win)
    273.  
    274.         out = out + lepe
    275.         # output linear
    276.         out = self.wo(out)
    277.  
    278.         # NOTE: use padding for semantic segmentation
    279.         # crop padded region
    280.         if self.auto_pad and (pad_r > 0 or pad_b > 0):
    281.             out = out[:, :H_in, :W_in, :].contiguous()
    282.  
    283.         if ret_attn_mask:
    284.             return out, r_weight, r_idx, attn_weight
    285.         else:
    286.             return rearrange(out, "n h w c -> n c h w")

    4.2修改二

    我们找到该文件'ultralytics/nn/tasks.py'在其中添加如下一行代码

    from ultralytics.nn.modules.Biformer import BiLevelRoutingAttention as Biformer

    添加完之后的效果如下图->

    4.3修改三

    这里需要注意体现出两种注意力机制的修改方式~

    4.2.1有参数的注意力机制修改

    现在我们已经将Biformer文件导入了模型中了，下一步我们就需要添加该机制到模型中让我们可以使用它，我们在步骤二的文件中''ultralytics/nn/tasks.py''按快捷键Ctrl+F可以进行文件搜索。

    当然如果你不想用快捷键也可以自己寻找大概在 650行左右，有一个方法的名字叫"parse_model"

    我们找到该方法对其进行修改，添加如下图所示内容。

    这里我们定义了一个字典，我们以后在想导入其它的注意力机制就可以重复步骤一和步骤二，然后在步骤三这里定义的字典中添加你导入的注意力机制名字即可。 

    4.2.2无参数的注意力机制修改

    无参数的注意力机制直接修改完步骤二就可以，直接跳过本步骤的修改直接进行配置注意力机制即可，无参数的注意力机制的修改三不用进行任何修改~

    4.4配置注意力机制

    恭喜你，到这里我们就已经成功的导入了注意力机制，离修改模型只差最后一步，我们需要找到如下文件进行修改"ultralytics/cfg/models/v8/yolov8.yaml",找到这个文件之后初始如下所示，

    我们可以在某一层中添加Biformer注意力机制，具体添加到哪里由你自己决定，我这里建议添加到  Neck层，也就是我们的特征融合层，添加之后的效果如下，这里我在三个地方添加了Biformer注意力机制。

    OK到此我们就添加了注意力机制到我们的模型里面了，下面我来讲一下添加的注意力机制中的参数是如何来的，

    • 首先-1这里我们不用管， 它代表上一个层的输入输入-1就是让模型自动帮我们算输入的大小！
    • 数字1代表这里我们的Biformer注意力机制执行一次
    • Biformer代表我们的注意力机制名字，本来类的名字不是这个我在前面导入的时候给他另命名了前面有讲到
    • [7，8]这里是根据Biformer定义的时候来的，你只需要输入前两个即可(需要注意的是无参数的注意力机制这里什么都不用填写可以看看你的无参数注意力机制需要什么那种超参数你给予赋值即可，不用从模型中获取任何的其它参数)。

    当然这两个参数你可以换，调成其它的试试效果。

    五、Neck部分

    持续更新~

    六、检测头

    持续更新~

    七、损失函数

    持续更新~

  • 相关阅读:
    学习用docker构建自己的镜像
    educoder_python：4-1-逻辑控制（if）第2关：求解一元二次方程组
    如何查看自己的GPU型号以及配置信息
    我梦想中的学习组织-勤学会
    分析大疆官网根据ip自动跳转到指定域名
    java后端怎么入门，自学还是培训
    2023年09月 Python（六级）真题解析#中国电子学会#全国青少年软件编程等级考试
    Spring Data JPA 项目配置与QueryDSL集成
    Linux三剑客：awk的高级用法
    Oracle集群管理-19C集群禁用numa和大页内存特性
  • 原文地址：https://blog.csdn.net/java1314777/article/details/134432710