学习资源站

YOLOv11改进-主干_Backbone篇-大核心卷积UniRepLknet目标检测网络（适配yolov11全系列轻量化）

一、本文介绍

本文给大家带来的改进机制是 特征提取网络 UniRepLknet ，该网络结构的重点在于使用Dilated Reparam Block和大核心指导原则，强调了 高效的 结构进行通道间通讯和空间聚合，以及使用带扩张的小核心进行重新参数化，该网络结构就是在LKNet基础上的一个升级版本， LKNet我们之前已经出过教程了 。 UniRepLknet 在 各种视觉任务中，包括图像分类、目标检测和语义分割 ，都显示出优异的性能。

欢迎大家订阅我的专栏一起学习YOLO！

（本文内容可根据yolov11的N、S、M、L、X进行二次缩放，轻量化更上一层）。

二、 UniRepLknet的框架原理

官方论文地址： 官方论文地址

官方代码地址： 官方代码地址

UniRepLKNet论文提出了一种新型的大核心卷积 神经网络 架构。这种架构通过结合非扩张小核心和扩张小核心层来增强非扩张大核心卷积层，旨在优化空间模式层次和表示能力。它强调了根据下游任务选择核心大小的重要性，并展示了该架构在图像识别以外领域（如音频、视频和时间序列数据）的通用性。此外，论文还展示了在各种任务上的领先性能，证明了其多功能性和高效性。

UniRepLKNet的主要创新点包括：

UniRepLKNet的结构上的创新点主要体现在其大核心卷积神经网络的设计上，这包括对大核心的高效利用，以及与传统ConvNets和变换器相比的独特构造。这种架构通过融合非扩张小核心和扩张小核心层来增强大核心层，优化了空间模式层次和网络的表示能力。此外，论文还提出了适用于大核心ConvNets的四个架构设计原则，旨在充分发挥大核心的独特优势，例如通过浅层结构观察更广阔视野，而不需要深入网络层次。

在图像中展示的UniRepLKNet架构设计中，一个显著的结构创新是LaRK（Large Kernel ）块，它包括本文提出的Dilated Reparam Block，一个SE（Squeeze-and-Excitation）块，前馈网络（FFN），以及批量归一化（BN）层。LaRK块与SmaK（Small Kernel）块的主要区别在于，LaRK使用深度分离的3x3卷积层代替了Dilated Reparam Block中的层。不同阶段的块通过步长为2的密集3x3卷积层实现的下采样块连接，而这些块可以在不同阶段灵活地排列。这种设计强调了结构的模块化和灵活性，以及通过大核心来增强 模型 的性能和效率。

在图2中，UniRepLKNet的Dilated Reparam Block通过使用扩张的小核心卷积层来增强非扩张的大核心层。这些扩张层从参数角度看等同于一个具有更大稀疏核心的非扩张卷积层，这使得整个块可以等效地转换成单个大核心卷积。通过重新 参数化 的过程，多个具有不同扩张率的小核心卷积层被合并成一个等效的大核心卷积层，从而在保持可学习参数数量和计算效率的同时，增强了网络对空间信息的捕获能力。这种设计创新为ConvNets提供了更广泛的感受野，而不会增加模型的深度。

三、 UniRepLknet的核心代码

下面的核心代码使用方式看章节四。 

  1. # UniRepLKNet: A Universal Perception Large-Kernel ConvNet for Audio, Video, Point Cloud, Time-Series and Image Recognition
  2. # Github source: https://github.com/AILab-CVC/UniRepLKNet
  3. # Licensed under The Apache License 2.0 License [see LICENSE for details]
  4. # Based on RepLKNet, ConvNeXt, timm, DINO and DeiT code bases
  5. # https://github.com/DingXiaoH/RepLKNet-pytorch
  6. # https://github.com/facebookresearch/ConvNeXt
  7. # https://github.com/rwightman/pytorch-image-models/tree/master/timm
  8. # https://github.com/facebookresearch/deit/
  9. # https://github.com/facebookresearch/dino
  10. # --------------------------------------------------------'
  11. import torch
  12. import torch.nn as nn
  13. import torch.nn.functional as F
  14. from timm.layers import trunc_normal_, DropPath, to_2tuple
  15. from functools import partial
  16. import torch.utils.checkpoint as checkpoint
  17. import numpy as np
  19. __all__ = ['unireplknet_a', 'unireplknet_f', 'unireplknet_p', 'unireplknet_n', 'unireplknet_t', 'unireplknet_s', 'unireplknet_b', 'unireplknet_l', 'unireplknet_xl']
  21. class GRNwithNHWC(nn.Module):
  22.     """ GRN (Global Response Normalization) layer
  23.     Originally proposed in ConvNeXt V2 (https://arxiv.org/abs/2301.00808)
  24.     This implementation is more efficient than the original (https://github.com/facebookresearch/ConvNeXt-V2)
  25.     We assume the inputs to this layer are (N, H, W, C)
  26.     """
  27.     def __init__(self, dim, use_bias=True):
  28.         super().__init__()
  29.         self.use_bias = use_bias
  30.         self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
  31.         if self.use_bias:
  32.             self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
  34.     def forward(self, x):
  35.         Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
  36.         Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
  37.         if self.use_bias:
  38.             return (self.gamma * Nx + 1) * x + self.beta
  39.         else:
  40.             return (self.gamma * Nx + 1) * x
  43. class NCHWtoNHWC(nn.Module):
  44.     def __init__(self):
  45.         super().__init__()
  47.     def forward(self, x):
  48.         return x.permute(0, 2, 3, 1)
  51. class NHWCtoNCHW(nn.Module):
  52.     def __init__(self):
  53.         super().__init__()
  55.     def forward(self, x):
  56.         return x.permute(0, 3, 1, 2)
  58. #================== This function decides which conv implementation (the native or iGEMM) to use
  59. #   Note that iGEMM large-kernel conv impl will be used if
  60. #       -   you attempt to do so (attempt_to_use_large_impl=True), and
  61. #       -   it has been installed (follow https://github.com/AILab-CVC/UniRepLKNet), and
  62. #       -   the conv layer is depth-wise, stride = 1, non-dilated, kernel_size > 5, and padding == kernel_size // 2
  63. def get_conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias,
  64.                attempt_use_lk_impl=True):
  65.     kernel_size = to_2tuple(kernel_size)
  66.     if padding is None:
  67.         padding = (kernel_size[0] // 2, kernel_size[1] // 2)
  68.     else:
  69.         padding = to_2tuple(padding)
  70.     need_large_impl = kernel_size[0] == kernel_size[1] and kernel_size[0] > 5 and padding == (kernel_size[0] // 2, kernel_size[1] // 2)
  72.     # if attempt_use_lk_impl and need_large_impl:
  73.     #     print('---------------- trying to import iGEMM implementation for large-kernel conv')
  74.     #     try:
  75.     #         from depthwise_conv2d_implicit_gemm import DepthWiseConv2dImplicitGEMM
  76.     #         print('---------------- found iGEMM implementation ')
  77.     #     except:
  78.     #         DepthWiseConv2dImplicitGEMM = None
  79.     #         print('---------------- found no iGEMM. use original conv. follow https://github.com/AILab-CVC/UniRepLKNet to install it.')
  80.     #     if DepthWiseConv2dImplicitGEMM is not None and need_large_impl and in_channels == out_channels \
  81.     #             and out_channels == groups and stride == 1 and dilation == 1:
  82.     #         print(f'===== iGEMM Efficient Conv Impl, channels {in_channels}, kernel size {kernel_size} =====')
  83.     #         return DepthWiseConv2dImplicitGEMM(in_channels, kernel_size, bias=bias)
  84.     return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
  85.                      padding=padding, dilation=dilation, groups=groups, bias=bias)
  88. def get_bn(dim, use_sync_bn=False):
  89.     if use_sync_bn:
  90.         return nn.SyncBatchNorm(dim)
  91.     else:
  92.         return nn.BatchNorm2d(dim)
  94. class SEBlock(nn.Module):
  95.     """
  96.     Squeeze-and-Excitation Block proposed in SENet (https://arxiv.org/abs/1709.01507)
  97.     We assume the inputs to this layer are (N, C, H, W)
  98.     """
  99.     def __init__(self, input_channels, internal_neurons):
  100.         super(SEBlock, self).__init__()
  101.         self.down = nn.Conv2d(in_channels=input_channels, out_channels=internal_neurons,
  102.                               kernel_size=1, stride=1, bias=True)
  103.         self.up = nn.Conv2d(in_channels=internal_neurons, out_channels=input_channels,
  104.                             kernel_size=1, stride=1, bias=True)
  105.         self.input_channels = input_channels
  106.         self.nonlinear = nn.ReLU(inplace=True)
  108.     def forward(self, inputs):
  109.         x = F.adaptive_avg_pool2d(inputs, output_size=(1, 1))
  110.         x = self.down(x)
  111.         x = self.nonlinear(x)
  112.         x = self.up(x)
  113.         x = F.sigmoid(x)
  114.         return inputs * x.view(-1, self.input_channels, 1, 1)
  116. def fuse_bn(conv, bn):
  117.     conv_bias = 0 if conv.bias is None else conv.bias
  118.     std = (bn.running_var + bn.eps).sqrt()
  119.     return conv.weight * (bn.weight / std).reshape(-1, 1, 1, 1), bn.bias + (conv_bias - bn.running_mean) * bn.weight / std
  121. def convert_dilated_to_nondilated(kernel, dilate_rate):
  122.     identity_kernel = torch.ones((1, 1, 1, 1)).to(kernel.device)
  123.     if kernel.size(1) == 1:
  124.         #   This is a DW kernel
  125.         dilated = F.conv_transpose2d(kernel, identity_kernel, stride=dilate_rate)
  126.         return dilated
  127.     else:
  128.         #   This is a dense or group-wise (but not DW) kernel
  129.         slices = []
  130.         for i in range(kernel.size(1)):
  131.             dilated = F.conv_transpose2d(kernel[:,i:i+1,:,:], identity_kernel, stride=dilate_rate)
  132.             slices.append(dilated)
  133.         return torch.cat(slices, dim=1)
  135. def merge_dilated_into_large_kernel(large_kernel, dilated_kernel, dilated_r):
  136.     large_k = large_kernel.size(2)
  137.     dilated_k = dilated_kernel.size(2)
  138.     equivalent_kernel_size = dilated_r * (dilated_k - 1) + 1
  139.     equivalent_kernel = convert_dilated_to_nondilated(dilated_kernel, dilated_r)
  140.     rows_to_pad = large_k // 2 - equivalent_kernel_size // 2
  141.     merged_kernel = large_kernel + F.pad(equivalent_kernel, [rows_to_pad] * 4)
  142.     return merged_kernel
  145. class DilatedReparamBlock(nn.Module):
  146.     """
  147.     Dilated Reparam Block proposed in UniRepLKNet (https://github.com/AILab-CVC/UniRepLKNet)
  148.     We assume the inputs to this block are (N, C, H, W)
  149.     """
  150.     def __init__(self, channels, kernel_size, deploy, use_sync_bn=False, attempt_use_lk_impl=True):
  151.         super().__init__()
  152.         self.lk_origin = get_conv2d(channels, channels, kernel_size, stride=1,
  153.                                     padding=kernel_size//2, dilation=1, groups=channels, bias=deploy,
  154.                                     attempt_use_lk_impl=attempt_use_lk_impl)
  155.         self.attempt_use_lk_impl = attempt_use_lk_impl
  157.         #   Default settings. We did not tune them carefully. Different settings may work better.
  158.         if kernel_size == 17:
  159.             self.kernel_sizes = [5, 9, 3, 3, 3]
  160.             self.dilates = [1, 2, 4, 5, 7]
  161.         elif kernel_size == 15:
  162.             self.kernel_sizes = [5, 7, 3, 3, 3]
  163.             self.dilates = [1, 2, 3, 5, 7]
  164.         elif kernel_size == 13:
  165.             self.kernel_sizes = [5, 7, 3, 3, 3]
  166.             self.dilates = [1, 2, 3, 4, 5]
  167.         elif kernel_size == 11:
  168.             self.kernel_sizes = [5, 5, 3, 3, 3]
  169.             self.dilates = [1, 2, 3, 4, 5]
  170.         elif kernel_size == 9:
  171.             self.kernel_sizes = [5, 5, 3, 3]
  172.             self.dilates = [1, 2, 3, 4]
  173.         elif kernel_size == 7:
  174.             self.kernel_sizes = [5, 3, 3]
  175.             self.dilates = [1, 2, 3]
  176.         elif kernel_size == 5:
  177.             self.kernel_sizes = [3, 3]
  178.             self.dilates = [1, 2]
  179.         else:
  180.             raise ValueError('Dilated Reparam Block requires kernel_size >= 5')
  182.         if not deploy:
  183.             self.origin_bn = get_bn(channels, use_sync_bn)
  184.             for k, r in zip(self.kernel_sizes, self.dilates):
  185.                 self.__setattr__('dil_conv_k{}_{}'.format(k, r),
  186.                                  nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=k, stride=1,
  187.                                            padding=(r * (k - 1) + 1) // 2, dilation=r, groups=channels,
  188.                                            bias=False))
  189.                 self.__setattr__('dil_bn_k{}_{}'.format(k, r), get_bn(channels, use_sync_bn=use_sync_bn))
  191.     def forward(self, x):
  192.         if not hasattr(self, 'origin_bn'):      # deploy mode
  193.             return self.lk_origin(x)
  194.         out = self.origin_bn(self.lk_origin(x))
  195.         for k, r in zip(self.kernel_sizes, self.dilates):
  196.             conv = self.__getattr__('dil_conv_k{}_{}'.format(k, r))
  197.             bn = self.__getattr__('dil_bn_k{}_{}'.format(k, r))
  198.             out = out + bn(conv(x))
  199.         return out
  201.     def merge_dilated_branches(self):
  202.         if hasattr(self, 'origin_bn'):
  203.             origin_k, origin_b = fuse_bn(self.lk_origin, self.origin_bn)
  204.             for k, r in zip(self.kernel_sizes, self.dilates):
  205.                 conv = self.__getattr__('dil_conv_k{}_{}'.format(k, r))
  206.                 bn = self.__getattr__('dil_bn_k{}_{}'.format(k, r))
  207.                 branch_k, branch_b = fuse_bn(conv, bn)
  208.                 origin_k = merge_dilated_into_large_kernel(origin_k, branch_k, r)
  209.                 origin_b += branch_b
  210.             merged_conv = get_conv2d(origin_k.size(0), origin_k.size(0), origin_k.size(2), stride=1,
  211.                                     padding=origin_k.size(2)//2, dilation=1, groups=origin_k.size(0), bias=True,
  212.                                     attempt_use_lk_impl=self.attempt_use_lk_impl)
  213.             merged_conv.weight.data = origin_k
  214.             merged_conv.bias.data = origin_b
  215.             self.lk_origin = merged_conv
  216.             self.__delattr__('origin_bn')
  217.             for k, r in zip(self.kernel_sizes, self.dilates):
  218.                 self.__delattr__('dil_conv_k{}_{}'.format(k, r))
  219.                 self.__delattr__('dil_bn_k{}_{}'.format(k, r))
  222. class UniRepLKNetBlock(nn.Module):
  224.     def __init__(self,
  225.                  dim,
  226.                  kernel_size,
  227.                  drop_path=0.,
  228.                  layer_scale_init_value=1e-6,
  229.                  deploy=False,
  230.                  attempt_use_lk_impl=True,
  231.                  with_cp=False,
  232.                  use_sync_bn=False,
  233.                  ffn_factor=4):
  234.         super().__init__()
  235.         self.with_cp = with_cp
  236.         # if deploy:
  237.         #     print('------------------------------- Note: deploy mode')
  238.         # if self.with_cp:
  239.         #     print('****** note with_cp = True, reduce memory consumption but may slow down training ******')
  241.         self.need_contiguous = (not deploy) or kernel_size >= 7
  243.         if kernel_size == 0:
  244.             self.dwconv = nn.Identity()
  245.             self.norm = nn.Identity()
  246.         elif deploy:
  247.             self.dwconv = get_conv2d(dim, dim, kernel_size=kernel_size, stride=1, padding=kernel_size // 2,
  248.                                      dilation=1, groups=dim, bias=True,
  249.                                      attempt_use_lk_impl=attempt_use_lk_impl)
  250.             self.norm = nn.Identity()
  251.         elif kernel_size >= 7:
  252.             self.dwconv = DilatedReparamBlock(dim, kernel_size, deploy=deploy,
  253.                                               use_sync_bn=use_sync_bn,
  254.                                               attempt_use_lk_impl=attempt_use_lk_impl)
  255.             self.norm = get_bn(dim, use_sync_bn=use_sync_bn)
  256.         elif kernel_size == 1:
  257.             self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, padding=kernel_size // 2,
  258.                                     dilation=1, groups=1, bias=deploy)
  259.             self.norm = get_bn(dim, use_sync_bn=use_sync_bn)
  260.         else:
  261.             assert kernel_size in [3, 5]
  262.             self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, stride=1, padding=kernel_size // 2,
  263.                                     dilation=1, groups=dim, bias=deploy)
  264.             self.norm = get_bn(dim, use_sync_bn=use_sync_bn)
  266.         self.se = SEBlock(dim, dim // 4)
  268.         ffn_dim = int(ffn_factor * dim)
  269.         self.pwconv1 = nn.Sequential(
  270.             NCHWtoNHWC(),
  271.             nn.Linear(dim, ffn_dim))
  272.         self.act = nn.Sequential(
  273.             nn.GELU(),
  274.             GRNwithNHWC(ffn_dim, use_bias=not deploy))
  275.         if deploy:
  276.             self.pwconv2 = nn.Sequential(
  277.                 nn.Linear(ffn_dim, dim),
  278.                 NHWCtoNCHW())
  279.         else:
  280.             self.pwconv2 = nn.Sequential(
  281.                 nn.Linear(ffn_dim, dim, bias=False),
  282.                 NHWCtoNCHW(),
  283.                 get_bn(dim, use_sync_bn=use_sync_bn))
  285.         self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim),
  286.                                   requires_grad=True) if (not deploy) and layer_scale_init_value is not None \
  287.                                                          and layer_scale_init_value > 0 else None
  288.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  290.     def forward(self, inputs):
  292.         def _f(x):
  293.             if self.need_contiguous:
  294.                 x = x.contiguous()
  295.             y = self.se(self.norm(self.dwconv(x)))
  296.             y = self.pwconv2(self.act(self.pwconv1(y)))
  297.             if self.gamma is not None:
  298.                 y = self.gamma.view(1, -1, 1, 1) * y
  299.             return self.drop_path(y) + x
  301.         if self.with_cp and inputs.requires_grad:
  302.             return checkpoint.checkpoint(_f, inputs)
  303.         else:
  304.             return _f(inputs)
  306.     def reparameterize(self):
  307.         if hasattr(self.dwconv, 'merge_dilated_branches'):
  308.             self.dwconv.merge_dilated_branches()
  309.         if hasattr(self.norm, 'running_var') and hasattr(self.dwconv, 'lk_origin'):
  310.             std = (self.norm.running_var + self.norm.eps).sqrt()
  311.             self.dwconv.lk_origin.weight.data *= (self.norm.weight / std).view(-1, 1, 1, 1)
  312.             self.dwconv.lk_origin.bias.data = self.norm.bias + (self.dwconv.lk_origin.bias - self.norm.running_mean) * self.norm.weight / std
  313.             self.norm = nn.Identity()
  314.         if self.gamma is not None:
  315.             final_scale = self.gamma.data
  316.             self.gamma = None
  317.         else:
  318.             final_scale = 1
  319.         if self.act[1].use_bias and len(self.pwconv2) == 3:
  320.             grn_bias = self.act[1].beta.data
  321.             self.act[1].__delattr__('beta')
  322.             self.act[1].use_bias = False
  323.             linear = self.pwconv2[0]
  324.             grn_bias_projected_bias = (linear.weight.data @ grn_bias.view(-1, 1)).squeeze()
  325.             bn = self.pwconv2[2]
  326.             std = (bn.running_var + bn.eps).sqrt()
  327.             new_linear = nn.Linear(linear.in_features, linear.out_features, bias=True)
  328.             new_linear.weight.data = linear.weight * (bn.weight / std * final_scale).view(-1, 1)
  329.             linear_bias = 0 if linear.bias is None else linear.bias.data
  330.             linear_bias += grn_bias_projected_bias
  331.             new_linear.bias.data = (bn.bias + (linear_bias - bn.running_mean) * bn.weight / std) * final_scale
  332.             self.pwconv2 = nn.Sequential(new_linear, self.pwconv2[1])
  336. default_UniRepLKNet_A_F_P_kernel_sizes = ((3, 3),
  337.                                       (13, 13),
  338.                                       (13, 13, 13, 13, 13, 13),
  339.                                       (13, 13))
  340. default_UniRepLKNet_N_kernel_sizes = ((3, 3),
  341.                                       (13, 13),
  342.                                       (13, 13, 13, 13, 13, 13, 13, 13),
  343.                                       (13, 13))
  344. default_UniRepLKNet_T_kernel_sizes = ((3, 3, 3),
  345.                                       (13, 13, 13),
  346.                                       (13, 3, 13, 3, 13, 3, 13, 3, 13, 3, 13, 3, 13, 3, 13, 3, 13, 3),
  347.                                       (13, 13, 13))
  348. default_UniRepLKNet_S_B_L_XL_kernel_sizes = ((3, 3, 3),
  349.                                              (13, 13, 13),
  350.                                              (13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3, 13, 3, 3),
  351.                                              (13, 13, 13))
  352. UniRepLKNet_A_F_P_depths = (2, 2, 6, 2)
  353. UniRepLKNet_N_depths = (2, 2, 8, 2)
  354. UniRepLKNet_T_depths = (3, 3, 18, 3)
  355. UniRepLKNet_S_B_L_XL_depths = (3, 3, 27, 3)
  357. default_depths_to_kernel_sizes = {
  358.     UniRepLKNet_A_F_P_depths: default_UniRepLKNet_A_F_P_kernel_sizes,
  359.     UniRepLKNet_N_depths: default_UniRepLKNet_N_kernel_sizes,
  360.     UniRepLKNet_T_depths: default_UniRepLKNet_T_kernel_sizes,
  361.     UniRepLKNet_S_B_L_XL_depths: default_UniRepLKNet_S_B_L_XL_kernel_sizes
  362. }
  364. class UniRepLKNet(nn.Module):
  365.     r""" UniRepLKNet
  366.         A PyTorch impl of UniRepLKNet
  368.     Args:
  369.         in_chans (int): Number of input image channels. Default: 3
  370.         num_classes (int): Number of classes for classification head. Default: 1000
  371.         depths (tuple(int)): Number of blocks at each stage. Default: (3, 3, 27, 3)
  372.         dims (int): Feature dimension at each stage. Default: (96, 192, 384, 768)
  373.         drop_path_rate (float): Stochastic depth rate. Default: 0.
  374.         layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
  375.         head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
  376.         kernel_sizes (tuple(tuple(int))): Kernel size for each block. None means using the default settings. Default: None.
  377.         deploy (bool): deploy = True means using the inference structure. Default: False
  378.         with_cp (bool): with_cp = True means using torch.utils.checkpoint to save GPU memory. Default: False
  379.         init_cfg (dict): weights to load. The easiest way to use UniRepLKNet with for OpenMMLab family. Default: None
  380.         attempt_use_lk_impl (bool): try to load the efficient iGEMM large-kernel impl. Setting it to False disabling the iGEMM impl. Default: True
  381.         use_sync_bn (bool): use_sync_bn = True means using sync BN. Use it if your batch size is small. Default: False
  382.     """
  383.     def __init__(self,
  384.                  in_chans=3,
  385.                  num_classes=1000,
  386.                  depths=(3, 3, 27, 3),
  387.                  dims=(96, 192, 384, 768),
  388.                  drop_path_rate=0.,
  389.                  layer_scale_init_value=1e-6,
  390.                  head_init_scale=1.,
  391.                  kernel_sizes=None,
  392.                  deploy=False,
  393.                  with_cp=False,
  394.                  init_cfg=None,
  395.                  attempt_use_lk_impl=True,
  396.                  use_sync_bn=False,
  397.                  **kwargs
  398.                  ):
  399.         super().__init__()
  401.         depths = tuple(depths)
  402.         if kernel_sizes is None:
  403.             if depths in default_depths_to_kernel_sizes:
  404.                 # print('=========== use default kernel size ')
  405.                 kernel_sizes = default_depths_to_kernel_sizes[depths]
  406.             else:
  407.                 raise ValueError('no default kernel size settings for the given depths, '
  408.                                  'please specify kernel sizes for each block, e.g., '
  409.                                  '((3, 3), (13, 13), (13, 13, 13, 13, 13, 13), (13, 13))')
  410.         # print(kernel_sizes)
  411.         for i in range(4):
  412.             assert len(kernel_sizes[i]) == depths[i], 'kernel sizes do not match the depths'
  414.         self.with_cp = with_cp
  416.         dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
  417.         # print('=========== drop path rates: ', dp_rates)
  419.         self.downsample_layers = nn.ModuleList()
  420.         self.downsample_layers.append(nn.Sequential(
  421.             nn.Conv2d(in_chans, dims[0] // 2, kernel_size=3, stride=2, padding=1),
  422.             LayerNorm(dims[0] // 2, eps=1e-6, data_format="channels_first"),
  423.             nn.GELU(),
  424.             nn.Conv2d(dims[0] // 2, dims[0], kernel_size=3, stride=2, padding=1),
  425.             LayerNorm(dims[0], eps=1e-6, data_format="channels_first")))
  427.         for i in range(3):
  428.             self.downsample_layers.append(nn.Sequential(
  429.                 nn.Conv2d(dims[i], dims[i + 1], kernel_size=3, stride=2, padding=1),
  430.                 LayerNorm(dims[i + 1], eps=1e-6, data_format="channels_first")))
  432.         self.stages = nn.ModuleList()
  434.         cur = 0
  435.         for i in range(4):
  436.             main_stage = nn.Sequential(
  437.                 *[UniRepLKNetBlock(dim=dims[i], kernel_size=kernel_sizes[i][j], drop_path=dp_rates[cur + j],
  438.                                    layer_scale_init_value=layer_scale_init_value, deploy=deploy,
  439.                                    attempt_use_lk_impl=attempt_use_lk_impl,
  440.                                    with_cp=with_cp, use_sync_bn=use_sync_bn) for j in
  441.                   range(depths[i])])
  442.             self.stages.append(main_stage)
  443.             cur += depths[i]
  445.         self.output_mode = 'features'
  446.         norm_layer = partial(LayerNorm, eps=1e-6, data_format="channels_first")
  447.         for i_layer in range(4):
  448.             layer = norm_layer(dims[i_layer])
  449.             layer_name = f'norm{i_layer}'
  450.             self.add_module(layer_name, layer)
  451.         self.width_list = [i.size(1) for i in self.forward(torch.randn(1, 3, 640, 640))]
  452.         self.apply(self._init_weights)
  454.     def _init_weights(self, m):
  455.         if isinstance(m, (nn.Conv2d, nn.Linear)):
  456.             trunc_normal_(m.weight, std=.02)
  457.             if hasattr(m, 'bias') and m.bias is not None:
  458.                 nn.init.constant_(m.bias, 0)
  460.     def forward(self, x):
  461.         if self.output_mode == 'logits':
  462.             for stage_idx in range(4):
  463.                 x = self.downsample_layers[stage_idx](x)
  464.                 x = self.stages[stage_idx](x)
  465.             x = self.norm(x.mean([-2, -1]))
  466.             x = self.head(x)
  467.             return x
  468.         elif self.output_mode == 'features':
  469.             outs = []
  470.             for stage_idx in range(4):
  471.                 x = self.downsample_layers[stage_idx](x)
  472.                 x = self.stages[stage_idx](x)
  473.                 outs.append(self.__getattr__(f'norm{stage_idx}')(x))
  474.             return outs
  475.         else:
  476.             raise ValueError('Defined new output mode?')
  478.     def switch_to_deploy(self):
  479.         for m in self.modules():
  480.             if hasattr(m, 'reparameterize'):
  481.                 m.reparameterize()
  485. class LayerNorm(nn.Module):
  486.     r""" LayerNorm implementation used in ConvNeXt
  487.     LayerNorm that supports two data formats: channels_last (default) or channels_first.
  488.     The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
  489.     shape (batch_size, height, width, channels) while channels_first corresponds to inputs
  490.     with shape (batch_size, channels, height, width).
  491.     """
  493.     def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last", reshape_last_to_first=False):
  494.         super().__init__()
  495.         self.weight = nn.Parameter(torch.ones(normalized_shape))
  496.         self.bias = nn.Parameter(torch.zeros(normalized_shape))
  497.         self.eps = eps
  498.         self.data_format = data_format
  499.         if self.data_format not in ["channels_last", "channels_first"]:
  500.             raise NotImplementedError
  501.         self.normalized_shape = (normalized_shape,)
  502.         self.reshape_last_to_first = reshape_last_to_first
  504.     def forward(self, x):
  505.         if self.data_format == "channels_last":
  506.             return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
  507.         elif self.data_format == "channels_first":
  508.             u = x.mean(1, keepdim=True)
  509.             s = (x - u).pow(2).mean(1, keepdim=True)
  510.             x = (x - u) / torch.sqrt(s + self.eps)
  511.             x = self.weight[:, None, None] * x + self.bias[:, None, None]
  512.             return x
  514. def update_weight(model_dict, weight_dict):
  515.     idx, temp_dict = 0, {}
  516.     for k, v in weight_dict.items():
  517.         if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
  518.             temp_dict[k] = v
  519.             idx += 1
  520.     model_dict.update(temp_dict)
  521.     print(f'loading weights... {idx}/{len(model_dict)} items')
  522.     return model_dict
  524. def unireplknet_a(pretrained='', **kwargs):
  525.     model = UniRepLKNet(depths=UniRepLKNet_A_F_P_depths, dims=(40, 80, 160, 320), **kwargs)
  526.     if pretrained:
  527.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  528.     return model
  530. def unireplknet_f(pretrained='', **kwargs):
  531.     model = UniRepLKNet(depths=UniRepLKNet_A_F_P_depths, dims=(48, 96, 192, 384), **kwargs)
  532.     if pretrained:
  533.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  534.     return model
  536. def unireplknet_p(pretrained='', **kwargs):
  537.     model = UniRepLKNet(depths=UniRepLKNet_A_F_P_depths, dims=(64, 128, 256, 512), **kwargs)
  538.     if pretrained:
  539.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  540.     return model
  542. def unireplknet_n(weights='', **kwargs):
  543.     model = UniRepLKNet(depths=UniRepLKNet_N_depths, dims=(80, 160, 320, 640), **kwargs)
  544.     if weights:
  545.         model.load_state_dict(update_weight(model.state_dict(), torch.load(weights)))
  546.     return model
  548. def unireplknet_t(pretrained='', **kwargs):
  549.     model = UniRepLKNet(depths=UniRepLKNet_T_depths, dims=(80, 160, 320, 640), **kwargs)
  550.     if pretrained:
  551.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  552.     return model
  554. def unireplknet_s(pretrained='', **kwargs):
  555.     model = UniRepLKNet(depths=UniRepLKNet_S_B_L_XL_depths, dims=(96, 192, 384, 768), **kwargs)
  556.     if pretrained:
  557.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  558.     return model
  560. def unireplknet_b(pretrained='', **kwargs):
  561.     model = UniRepLKNet(depths=UniRepLKNet_S_B_L_XL_depths, dims=(128, 256, 512, 1024), **kwargs)
  562.     if pretrained:
  563.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  564.     return model
  566. def unireplknet_l(pretrained='', **kwargs):
  567.     model = UniRepLKNet(depths=UniRepLKNet_S_B_L_XL_depths, dims=(192, 384, 768, 1536), **kwargs)
  568.     if pretrained:
  569.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  570.     return model
  572. def unireplknet_xl(pretrained='', **kwargs):
  573.     model = UniRepLKNet(depths=UniRepLKNet_S_B_L_XL_depths, dims=(256, 512, 1024, 2048), **kwargs)
  574.     if pretrained:
  575.         model.load_state_dict(update_weight(model.state_dict(), torch.load(pretrained)))
  576.     return model

四、手把手教你添加UniRepLknet机制

4.1 修改一

第一步还是建立文件，我们找到如下ultralytics/nn文件夹下建立一个目录名字呢就是'Addmodules'文件夹( 用群内的文件的话已经有了无需新建) ！然后在其内部建立一个新的py文件将核心代码复制粘贴进去即可

4.2 修改二

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

4.3 修改三

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

从今天开始以后的教程就都统一成这个样子了，因为我默认大家用了我群内的文件来进行修改！！

4.4 修改四

添加如下两行代码！！！

​

4.5 修改五

找到七百多行大概把具体看图片，按照图片来修改就行，添加红框内的部分，注意没有()只是 函数 名。 

  1.         elif m in {自行添加对应的模型即可，下面都是一样的}:
  2.             m = m(*args)
  3.             c2 = m.width_list  # 返回通道列表
  4.             backbone = True

4.6 修改六

下面的两个红框内都是需要改动的。

​ 

  1.         if isinstance(c2, list):
  2.             m_ = m
  3.             m_.backbone = True
  4.         else:
  5.             m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
  6.             t = str(m)[8:-2].replace('__main__.', '')  # module type
  9.         m.np = sum(x.numel() for x in m_.parameters())  # number params
  10.         m_.i, m_.f, m_.type = i + 4 if backbone else i, f, t  # attach index, 'from' index, type

4.7 修改七

如下的也需要修改，全部按照我的来。

​

代码如下把原先的代码替换了即可。 

  1.         if verbose:
  2.             LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f}  {t:<45}{str(args):<30}')  # print
  4.         save.extend(x % (i + 4 if backbone else i) for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
  5.         layers.append(m_)
  6.         if i == 0:
  7.             ch = []
  8.         if isinstance(c2, list):
  9.             ch.extend(c2)
  10.             if len(c2) != 5:
  11.                 ch.insert(0, 0)
  12.         else:
  13.             ch.append(c2)

4.8 修改八

修改八和前面的都不太一样，需要修改前向传播中的一个部分， 已经离开了parse_model方法了。

可以在图片中开代码行数，没有离开task.py文件都是同一个文件。 同时这个部分有好几个前向传播都很相似，大家不要看错了， 是70多行左右的！！！，同时我后面提供了代码，大家直接复制粘贴即可，有时间我针对这里会出一个视频。

​​

代码如下-> 

  1.     def _predict_once(self, x, profile=False, visualize=False, embed=None):
  2.         """
  3.         Perform a forward pass through the network.
  5.         Args:
  6.             x (torch.Tensor): The input tensor to the model.
  7.             profile (bool):  Print the computation time of each layer if True, defaults to False.
  8.             visualize (bool): Save the feature maps of the model if True, defaults to False.
  9.             embed (list, optional): A list of feature vectors/embeddings to return.
  11.         Returns:
  12.             (torch.Tensor): The last output of the model.
  13.         """
  14.         y, dt, embeddings = [], [], []  # outputs
  15.         for m in self.model:
  16.             if m.f != -1:  # if not from previous layer
  17.                 x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers
  18.             if profile:
  19.                 self._profile_one_layer(m, x, dt)
  20.             if hasattr(m, 'backbone'):
  21.                 x = m(x)
  22.                 if len(x) != 5:  # 0 - 5
  23.                     x.insert(0, None)
  24.                 for index, i in enumerate(x):
  25.                     if index in self.save:
  26.                         y.append(i)
  27.                     else:
  28.                         y.append(None)
  29.                 x = x[-1]  # 最后一个输出传给下一层
  30.             else:
  31.                 x = m(x)  # run
  32.                 y.append(x if m.i in self.save else None)  # save output
  33.             if visualize:
  34.                 feature_visualization(x, m.type, m.i, save_dir=visualize)
  35.             if embed and m.i in embed:
  36.                 embeddings.append(nn.functional.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1))  # flatten
  37.                 if m.i == max(embed):
  38.                     return torch.unbind(torch.cat(embeddings, 1), dim=0)
  39.         return x

到这里就完成了修改部分，但是这里面细节很多，大家千万要注意不要替换多余的代码，导致报错，也不要拉下任何一部，都会导致运行失败，而且报错很难排查！！！很难排查！！！

注意！！！ 额外的修改！

关注我的其实都知道，我大部分的修改都是一样的，这个网络需要额外的修改一步，就是s一个参数，将下面的s改为640！！！即可完美运行！！

打印计算量问题解决方案

我们找到如下文件'ultralytics/utils/torch_utils.py'按照如下的图片进行修改，否则容易打印不出来计算量。

​

注意事项！！！

如果大家在验证的时候报错形状不匹配的错误可以固定 验证集 的图片尺寸，方法如下 ->

找到下面这个文件ultralytics/ models /yolo/detect/train.py然后其中有一个类是DetectionTrainer class中的build_dataset函数中的一个参数rect=mode == 'val'改为rect=False

五、UniRepLknet的yaml文件

复制如下yaml文件进行运行！！！

5.1 UniRepLknet 的yaml文件版本1

此版本训练信息：YOLO11-UniRepLKNet summary: 570 layers, 1,876,231 parameters, 1,876,215 gradients, 4.5 GFLOPs

使用说明：# 下面 [-1, 1, UniRepLKNet, [0.25]] 参数位置的0.25是通道放缩的系数, YOLOv11N是0.25 YOLOv11S是0.5 YOLOv11M是1. YOLOv11l是1 YOLOv11是1.5大家根据自己训练的YOLO版本设定即可.

# 本文支持版本有'unireplknet_a', 'unireplknet_f', 'unireplknet_p', 'unireplknet_n', 'unireplknet_t', 'unireplknet_s', 'unireplknet_b', 'unireplknet_l', 'unireplknet_xl' 

  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. # 下面 [-1, 1, LSKNet, [0.25，0.5]] 参数位置的0.25是通道放缩的系数, YOLOv11N是0.25 YOLOv11S是0.5 YOLOv11M是1. YOLOv11l是1 YOLOv11是1.5大家根据自己训练的YOLO版本设定即可.
  15. # 支持的版本有： __all__ = ['unireplknet_a', 'unireplknet_f', 'unireplknet_p', 'unireplknet_n', 'unireplknet_t', 'unireplknet_s', 'unireplknet_b', 'unireplknet_l', 'unireplknet_xl']
  16. # YOLO11n backbone
  17. backbone:
  18.   # [from, repeats, module, args]
  19.   - [-1, 1, unireplknet_a, [0.25]] # 0-4 P1/2
  20.   - [-1, 1, SPPF, [1024, 5]] # 5
  21.   - [-1, 2, C2PSA, [1024]] # 6
  23. # YOLO11n head
  24. head:
  25.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  26.   - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  27.   - [-1, 2, C3k2, [512, False]] # 9
  29.   - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  30.   - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  31.   - [-1, 2, C3k2, [256, False]] # 12 (P3/8-small)
  33.   - [-1, 1, Conv, [256, 3, 2]]
  34.   - [[-1, 9], 1, Concat, [1]] # cat head P4
  35.   - [-1, 2, C3k2, [512, False]] # 15 (P4/16-medium)
  37.   - [-1, 1, Conv, [512, 3, 2]]
  38.   - [[-1, 6], 1, Concat, [1]] # cat head P5
  39.   - [-1, 2, C3k2, [1024, True]] # 18 (P5/32-large)
  41.   - [[12, 15, 18], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 训练文件 

  1. import warnings
  2. warnings.filterwarnings('ignore')
  3. from ultralytics import YOLO
  5. if __name__ == '__main__':
  6.     model = YOLO('ultralytics/cfg/models/v8/yolov8-C2f-FasterBlock.yaml')
  7.     # model.load('yolov8n.pt') # loading pretrain weights
  8.     model.train(data=r'替换数据集yaml文件地址',
  9.                 # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose
  10.                 cache=False,
  11.                 imgsz=640,
  12.                 epochs=150,
  13.                 single_cls=False,  # 是否是单类别检测
  14.                 batch=4,
  15.                 close_mosaic=10,
  16.                 workers=0,
  17.                 device='0',
  18.                 optimizer='SGD', # using SGD
  19.                 # resume='', # 如过想续训就设置last.pt的地址
  20.                 amp=False,  # 如果出现训练损失为Nan可以关闭amp
  21.                 project='runs/train',
  22.                 name='exp',
  23.                 )

六、成功运行记录

下面是成功运行的截图，已经完成了有1个epochs的训练，图片太大截不全第2个epochs，这里改完之后打印出了点问题，但是不影响任何功能，后期我找时间修复一下这个问题。

​​

七、本文总结

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

​​​