学习资源站

29-添加DSConv注意力卷积（ICCV 2023_用于管状结构分割的动态蛇形卷积）_dsconv是哪篇论文提出的

YOLOv5改进系列（28）——添加DSConv注意力卷积（ICCV 2023|用于管状结构分割的动态蛇形卷积）

🚀 一、DSConv介绍 

学习资料：

1.1 DSConv简介

背景

管状结构（例如血管、道路）是临床、自然界等各领域场景中十分重要的一种结构，其精确分割可以保证下游任务的准确性与效率。但管状结构的精确提取仍然面临着众多挑战：

  • 细长且脆弱的局部结构。如下图所示，细长的结构仅占整个图像的一小部分，像素的组成有限。此外，这些结构容易受到复杂背景的干扰，因此模型很难精确分辨目标的细微变化，从而导致分割出现破碎与断裂。
  • 复杂且多变的全局形态。如下图所示，我们可以看出细长管状结构复杂多变的形态，即使在同一张图像中也是如此。位于不同区域的目标的形态变化取决于分支的数量、分叉的位置，路径长度以及其在图像中的位置。因此当数据表现出未曾见过的形态特征时，模型倾向于过拟合到已见过的特征，无法识别未见过的特征形态，从而导致泛化性较弱。

本文主要工作 

本文关注到管状结构细长连续的特点，并利用这一信息在神经网络以下三个阶段同时增强感知：特征提取、特征融合和损失约束。分别设计了动态蛇形卷积（Dynamic Snake Convolution），多视角特征融合策略与连续性拓扑约束损失。我们同时给出了基于 2D 和 3D 的方法设计，通过实验证明了本文所提出的 DSCNet 在管状结构分割任务上提供了更好的精度和连续性。 

 1.2 动态蛇形卷积

目的：

  • 希望卷积核一方面能够自由地贴合结构学习特征
  • 另一方面能够在约束条件下不偏离目标结构太远

可变形卷积：

  • 操控单个卷积核形变的所有偏置(offset)，是在网络中一次性全部学到的
  • 对于这一个偏置只有一个范围的约束，即感受野范围(extend)
  • 控制所有的卷积发生形变，是依赖于整个网络最终的损失约束回传，这个变化过程是相当自由的。 

1.3  多视角特征融合策略

目的：

  • 管状结构的走向与视角不是单一的，因此在设计中融合多视角特征也是必然的选择。

挑战：

  • 融合更多的特征会导致更大的网络负载以及出现冗余。

方法：

  • 在特征融合的训练过程中加入了分组与dropout的策略，一定程度上缓解了网络内内存的压力并避免模型陷入过拟合。 

 1.4 连续性拓扑约束损失

目的：

  • 构建数据的拓扑结构，并提取复杂管状结构中的高维关系，也就是持续同源性（Persistence Homology, PH）。

启发：

  • 假设 PO 的上端存在着一个异常的离散点（横坐标表示出现的时间，纵坐标表示消失的时间），这表明存在一个构件直到最后才与其他构件获得连接从而消失。

方法： 

  • 本文中采用的是豪斯多夫距离（Hausdorff Distance, HD），HD 也是用于衡量点集相似度的一个重要算法，对离散点也非常敏感。
  1. # -*- coding: utf-8 -*-
  3. import torch
  5. from torch import nn
  6. from torch.nn.functional import max_pool3d
  9. class crossentry(nn.Module):
  11.     def __init__(self):
  12.         super().__init__()
  14.     def forward(self, y_true, y_pred):
  15.         smooth = 1e-6
  16.         return -torch.mean(y_true * torch.log(y_pred + smooth))
  19. class cross_loss(nn.Module):
  21.     def __init__(self):
  22.         super().__init__()
  24.     def forward(self, y_true, y_pred):
  25.         smooth = 1e-6
  26.         return -torch.mean(y_true * torch.log(y_pred + smooth) +
  27.                            (1 - y_true) * torch.log(1 - y_pred + smooth))
  30. '''
  31. Another Loss Function proposed by us in IEEE transactions on Image Precessing:
  32. Paper: https://ieeexplore.ieee.org/abstract/document/9611074
  33. Code: https://github.com/YaoleiQi/Examinee-Examiner-Network
  34. '''
  37. class Dropoutput_Layer(nn.Module):
  39.     def __init__(self):
  40.         super().__init__()
  42.     def forward(self, y_true, y_pred, alpha=0.4):
  43.         smooth = 1e-6
  44.         w = torch.abs(y_true - y_pred)
  45.         w = torch.round(w + alpha)
  46.         loss_ce = (
  47.             -((torch.sum(w * y_true * torch.log(y_pred + smooth)) /
  48.                torch.sum(w * y_true + smooth)) +
  49.               (torch.sum(w * (1 - y_true) * torch.log(1 - y_pred + smooth)) /
  50.                torch.sum(w * (1 - y_true) + smooth))) / 2)
  51.         return loss_ce

🚀二、具体添加方法

2.1 添加顺序 

（1）models/common.py    -->  加入新增的网络结构

（2）     models/yolo.py       -->  设定网络结构的传参细节，将DSConv类名加入其中。（当新的自定义模块中存在输入输出维度时，要使用qw调整输出维度）
（3） models/yolov5*.yaml  -->  新建一个文件夹，如yolov5s_DSConv.yaml，修改现有模型结构配置文件。（当引入新的层时，要修改后续的结构中的from参数）
（4）         train.py                -->  修改‘--cfg’默认参数，训练时指定模型结构配置文件 

2.2 具体添加步骤 

第①步：在common.py中添加DCConv模块 

将下面的DSConv代码复制粘贴到common.py文件的末尾。

  1. # by:迪菲赫尔曼
  2. import warnings
  3. import torch
  4. from torch import nn
  6. warnings.filterwarnings("ignore")
  8. """
  9. This code is mainly the deformation process of our DSConv
  10. """
  12. class DSConv(nn.Module):
  14.     def __init__(self, in_ch, out_ch, kernel_size, extend_scope, morph,
  15.                  if_offset):
  16.         """
  17.         动态蛇形卷积
  18.         :param in_ch: 输入通道
  19.         :param out_ch: 输出通道
  20.         :param kernel_size: 卷积核的大小
  21.         :param extend_scope: 扩展范围（默认为此方法的1）
  22.         :param morph: 卷积核的形态主要分为两种类型，沿x轴（0）和沿y轴（1）（详细信息请参阅论文）
  23.         :param if_offset: 是否需要变形，如果为False，则是标准卷积核
  24.         """
  25.         super(DSConv, self).__init__()
  26.         # use the <offset_conv> to learn the deformable offset
  27.         # offset_conv: 学习可变形偏移的卷积层
  28.         self.offset_conv = nn.Conv2d(in_ch, 2 * kernel_size, 3, padding=1)
  29.         self.bn = nn.BatchNorm2d(2 * kernel_size)
  30.         self.kernel_size = kernel_size
  32.         # two types of the DSConv (along x-axis and y-axis)
  33.         # dsc_conv_x 和 dsc_conv_y：两种动态蛇形卷积层，分别沿x轴和y轴。
  34.         self.dsc_conv_x = nn.Conv2d(
  35.             in_ch,
  36.             out_ch,
  37.             kernel_size=(kernel_size, 1),
  38.             stride=(kernel_size, 1),
  39.             padding=0,
  40.         )
  41.         self.dsc_conv_y = nn.Conv2d(
  42.             in_ch,
  43.             out_ch,
  44.             kernel_size=(1, kernel_size),
  45.             stride=(1, kernel_size),
  46.             padding=0,
  47.         )
  49.         # gn：组归一化层
  50.         self.gn = nn.GroupNorm(out_ch // 4, out_ch)
  51.         self.relu = nn.ReLU(inplace=True)
  53.         # extend_scope：扩展范围
  54.         self.extend_scope = extend_scope
  55.         # morph：卷积核形态的类型
  56.         self.morph = morph
  57.         # if_offset：指示是否需要变形的布尔值
  58.         self.if_offset = if_offset
  60.     def forward(self, f):
  61.         offset = self.offset_conv(f)
  62.         offset = self.bn(offset)
  63.         # We need a range of deformation between -1 and 1 to mimic the snake's swing
  64.         offset = torch.tanh(offset)
  65.         input_shape = f.shape
  66.         dsc = DSC(input_shape, self.kernel_size, self.extend_scope, self.morph)
  67.         deformed_feature = dsc.deform_conv(f, offset, self.if_offset)
  68.         if self.morph == 0:
  69.             x = self.dsc_conv_x(deformed_feature.type(f.dtype))
  70.             x = self.gn(x)
  71.             x = self.relu(x)
  72.             return x
  73.         else:
  74.             x = self.dsc_conv_y(deformed_feature.type(f.dtype))
  75.             x = self.gn(x)
  76.             x = self.relu(x)
  77.             return x
  80. # Core code, for ease of understanding, we mark the dimensions of input and output next to the code
  81. class DSC(object):
  83.     def __init__(self, input_shape, kernel_size, extend_scope, morph):
  84.         self.num_points = kernel_size
  85.         self.width = input_shape[2]
  86.         self.height = input_shape[3]
  87.         self.morph = morph
  88.         self.extend_scope = extend_scope  # offset (-1 ~ 1) * extend_scope
  90.         # define feature map shape
  91.         """
  92.         B: Batch size  C: Channel  W: Width  H: Height
  93.         """
  94.         self.num_batch = input_shape[0]
  95.         self.num_channels = input_shape[1]
  97.     """
  98.     input: offset [B,2*K,W,H]  K: Kernel size (2*K: 2D image, deformation contains <x_offset> and <y_offset>)
  99.     output_x: [B,1,W,K*H]   coordinate map
  100.     output_y: [B,1,K*W,H]   coordinate map
  101.     """
  103.     def _coordinate_map_3D(self, offset, if_offset):
  104.         """
  105.         1.输入为偏移 (offset) 和是否需要偏移 (if_offset)。
  106.         2.根据输入特征图的形状、卷积核大小、扩展范围以及形态类型，生成二维坐标映射。
  107.         3.如果形态类型为0，表示沿x轴，生成y坐标映射；如果形态类型为1，表示沿y轴，生成x坐标映射。
  108.         4.根据偏移和扩展范围调整坐标映射。
  109.         5.返回生成的坐标映射。
  110.         """
  111.         device = offset.device
  112.         # offset
  113.         y_offset, x_offset = torch.split(offset, self.num_points, dim=1)
  115.         y_center = torch.arange(0, self.width).repeat([self.height])
  116.         y_center = y_center.reshape(self.height, self.width)
  117.         y_center = y_center.permute(1, 0)
  118.         y_center = y_center.reshape([-1, self.width, self.height])
  119.         y_center = y_center.repeat([self.num_points, 1, 1]).float()
  120.         y_center = y_center.unsqueeze(0)
  122.         x_center = torch.arange(0, self.height).repeat([self.width])
  123.         x_center = x_center.reshape(self.width, self.height)
  124.         x_center = x_center.permute(0, 1)
  125.         x_center = x_center.reshape([-1, self.width, self.height])
  126.         x_center = x_center.repeat([self.num_points, 1, 1]).float()
  127.         x_center = x_center.unsqueeze(0)
  129.         if self.morph == 0:
  130.             """
  131.             Initialize the kernel and flatten the kernel
  132.                 y: only need 0
  133.                 x: -num_points//2 ~ num_points//2 (Determined by the kernel size)
  134.                 !!! The related PPT will be submitted later, and the PPT will contain the whole changes of each step
  135.             """
  136.             y = torch.linspace(0, 0, 1)
  137.             x = torch.linspace(
  138.                 -int(self.num_points // 2),
  139.                 int(self.num_points // 2),
  140.                 int(self.num_points),
  141.             )
  143.             y, x = torch.meshgrid(y, x)
  144.             y_spread = y.reshape(-1, 1)
  145.             x_spread = x.reshape(-1, 1)
  147.             y_grid = y_spread.repeat([1, self.width * self.height])
  148.             y_grid = y_grid.reshape([self.num_points, self.width, self.height])
  149.             y_grid = y_grid.unsqueeze(0)  # [B*K*K, W,H]
  151.             x_grid = x_spread.repeat([1, self.width * self.height])
  152.             x_grid = x_grid.reshape([self.num_points, self.width, self.height])
  153.             x_grid = x_grid.unsqueeze(0)  # [B*K*K, W,H]
  155.             y_new = y_center + y_grid
  156.             x_new = x_center + x_grid
  158.             y_new = y_new.repeat(self.num_batch, 1, 1, 1).to(device)
  159.             x_new = x_new.repeat(self.num_batch, 1, 1, 1).to(device)
  161.             y_offset_new = y_offset.detach().clone()
  163.             if if_offset:
  164.                 y_offset = y_offset.permute(1, 0, 2, 3)
  165.                 y_offset_new = y_offset_new.permute(1, 0, 2, 3)
  166.                 center = int(self.num_points // 2)
  168.                 # The center position remains unchanged and the rest of the positions begin to swing
  169.                 # This part is quite simple. The main idea is that "offset is an iterative process"
  170.                 y_offset_new[center] = 0
  171.                 for index in range(1, center):
  172.                     y_offset_new[center + index] = (y_offset_new[center + index - 1] + y_offset[center + index])
  173.                     y_offset_new[center - index] = (y_offset_new[center - index + 1] + y_offset[center - index])
  174.                 y_offset_new = y_offset_new.permute(1, 0, 2, 3).to(device)
  175.                 y_new = y_new.add(y_offset_new.mul(self.extend_scope))
  177.             y_new = y_new.reshape(
  178.                 [self.num_batch, self.num_points, 1, self.width, self.height])
  179.             y_new = y_new.permute(0, 3, 1, 4, 2)
  180.             y_new = y_new.reshape([
  181.                 self.num_batch, self.num_points * self.width, 1 * self.height
  182.             ])
  183.             x_new = x_new.reshape(
  184.                 [self.num_batch, self.num_points, 1, self.width, self.height])
  185.             x_new = x_new.permute(0, 3, 1, 4, 2)
  186.             x_new = x_new.reshape([
  187.                 self.num_batch, self.num_points * self.width, 1 * self.height
  188.             ])
  189.             return y_new, x_new
  191.         else:
  192.             """
  193.             Initialize the kernel and flatten the kernel
  194.                 y: -num_points//2 ~ num_points//2 (Determined by the kernel size)
  195.                 x: only need 0
  196.             """
  197.             y = torch.linspace(
  198.                 -int(self.num_points // 2),
  199.                 int(self.num_points // 2),
  200.                 int(self.num_points),
  201.             )
  202.             x = torch.linspace(0, 0, 1)
  204.             y, x = torch.meshgrid(y, x)
  205.             y_spread = y.reshape(-1, 1)
  206.             x_spread = x.reshape(-1, 1)
  208.             y_grid = y_spread.repeat([1, self.width * self.height])
  209.             y_grid = y_grid.reshape([self.num_points, self.width, self.height])
  210.             y_grid = y_grid.unsqueeze(0)
  212.             x_grid = x_spread.repeat([1, self.width * self.height])
  213.             x_grid = x_grid.reshape([self.num_points, self.width, self.height])
  214.             x_grid = x_grid.unsqueeze(0)
  216.             y_new = y_center + y_grid
  217.             x_new = x_center + x_grid
  219.             y_new = y_new.repeat(self.num_batch, 1, 1, 1)
  220.             x_new = x_new.repeat(self.num_batch, 1, 1, 1)
  222.             y_new = y_new.to(device)
  223.             x_new = x_new.to(device)
  224.             x_offset_new = x_offset.detach().clone()
  226.             if if_offset:
  227.                 x_offset = x_offset.permute(1, 0, 2, 3)
  228.                 x_offset_new = x_offset_new.permute(1, 0, 2, 3)
  229.                 center = int(self.num_points // 2)
  230.                 x_offset_new[center] = 0
  231.                 for index in range(1, center):
  232.                     x_offset_new[center + index] = (x_offset_new[center + index - 1] + x_offset[center + index])
  233.                     x_offset_new[center - index] = (x_offset_new[center - index + 1] + x_offset[center - index])
  234.                 x_offset_new = x_offset_new.permute(1, 0, 2, 3).to(device)
  235.                 x_new = x_new.add(x_offset_new.mul(self.extend_scope))
  237.             y_new = y_new.reshape(
  238.                 [self.num_batch, 1, self.num_points, self.width, self.height])
  239.             y_new = y_new.permute(0, 3, 1, 4, 2)
  240.             y_new = y_new.reshape([
  241.                 self.num_batch, 1 * self.width, self.num_points * self.height
  242.             ])
  243.             x_new = x_new.reshape(
  244.                 [self.num_batch, 1, self.num_points, self.width, self.height])
  245.             x_new = x_new.permute(0, 3, 1, 4, 2)
  246.             x_new = x_new.reshape([
  247.                 self.num_batch, 1 * self.width, self.num_points * self.height
  248.             ])
  249.             return y_new, x_new
  251.     """
  252.     input: input feature map [N,C,D,W,H]；coordinate map [N,K*D,K*W,K*H] 
  253.     output: [N,1,K*D,K*W,K*H]  deformed feature map
  254.     """
  256.     def _bilinear_interpolate_3D(self, input_feature, y, x):
  257.         """
  258.          1.输入为输入特征图 (input_feature)、y坐标映射 (y) 和x坐标映射 (x)。
  259.          2.进行三维双线性插值，获取变形后的特征。
  260.          3.返回插值得到的变形特征。
  261.         """
  262.         device = input_feature.device
  263.         y = y.reshape([-1]).float()
  264.         x = x.reshape([-1]).float()
  266.         zero = torch.zeros([]).int()
  267.         max_y = self.width - 1
  268.         max_x = self.height - 1
  270.         # find 8 grid locations
  271.         y0 = torch.floor(y).int()
  272.         y1 = y0 + 1
  273.         x0 = torch.floor(x).int()
  274.         x1 = x0 + 1
  276.         # clip out coordinates exceeding feature map volume
  277.         y0 = torch.clamp(y0, zero, max_y)
  278.         y1 = torch.clamp(y1, zero, max_y)
  279.         x0 = torch.clamp(x0, zero, max_x)
  280.         x1 = torch.clamp(x1, zero, max_x)
  282.         input_feature_flat = input_feature.flatten()
  283.         input_feature_flat = input_feature_flat.reshape(
  284.             self.num_batch, self.num_channels, self.width, self.height)
  285.         input_feature_flat = input_feature_flat.permute(0, 2, 3, 1)
  286.         input_feature_flat = input_feature_flat.reshape(-1, self.num_channels)
  287.         dimension = self.height * self.width
  289.         base = torch.arange(self.num_batch) * dimension
  290.         base = base.reshape([-1, 1]).float()
  292.         repeat = torch.ones([self.num_points * self.width * self.height
  293.                              ]).unsqueeze(0)
  294.         repeat = repeat.float()
  296.         base = torch.matmul(base, repeat)
  297.         base = base.reshape([-1])
  299.         base = base.to(device)
  301.         base_y0 = base + y0 * self.height
  302.         base_y1 = base + y1 * self.height
  304.         # top rectangle of the neighbourhood volume
  305.         index_a0 = base_y0 - base + x0
  306.         index_c0 = base_y0 - base + x1
  308.         # bottom rectangle of the neighbourhood volume
  309.         index_a1 = base_y1 - base + x0
  310.         index_c1 = base_y1 - base + x1
  312.         # get 8 grid values
  313.         value_a0 = input_feature_flat[index_a0.type(torch.int64)].to(device)
  314.         value_c0 = input_feature_flat[index_c0.type(torch.int64)].to(device)
  315.         value_a1 = input_feature_flat[index_a1.type(torch.int64)].to(device)
  316.         value_c1 = input_feature_flat[index_c1.type(torch.int64)].to(device)
  318.         # find 8 grid locations
  319.         y0 = torch.floor(y).int()
  320.         y1 = y0 + 1
  321.         x0 = torch.floor(x).int()
  322.         x1 = x0 + 1
  324.         # clip out coordinates exceeding feature map volume
  325.         y0 = torch.clamp(y0, zero, max_y + 1)
  326.         y1 = torch.clamp(y1, zero, max_y + 1)
  327.         x0 = torch.clamp(x0, zero, max_x + 1)
  328.         x1 = torch.clamp(x1, zero, max_x + 1)
  330.         x0_float = x0.float()
  331.         x1_float = x1.float()
  332.         y0_float = y0.float()
  333.         y1_float = y1.float()
  335.         vol_a0 = ((y1_float - y) * (x1_float - x)).unsqueeze(-1).to(device)
  336.         vol_c0 = ((y1_float - y) * (x - x0_float)).unsqueeze(-1).to(device)
  337.         vol_a1 = ((y - y0_float) * (x1_float - x)).unsqueeze(-1).to(device)
  338.         vol_c1 = ((y - y0_float) * (x - x0_float)).unsqueeze(-1).to(device)
  340.         outputs = (value_a0 * vol_a0 + value_c0 * vol_c0 + value_a1 * vol_a1 +
  341.                    value_c1 * vol_c1)
  343.         if self.morph == 0:
  344.             outputs = outputs.reshape([
  345.                 self.num_batch,
  346.                 self.num_points * self.width,
  347.                 1 * self.height,
  348.                 self.num_channels,
  349.             ])
  350.             outputs = outputs.permute(0, 3, 1, 2)
  351.         else:
  352.             outputs = outputs.reshape([
  353.                 self.num_batch,
  354.                 1 * self.width,
  355.                 self.num_points * self.height,
  356.                 self.num_channels,
  357.             ])
  358.             outputs = outputs.permute(0, 3, 1, 2)
  359.         return outputs
  361.     def deform_conv(self, input, offset, if_offset):
  362.         """
  363.         1.输入为原始特征图 (input)、偏移 (offset) 和是否需要偏移 (if_offset)。
  364.         2.调用 _coordinate_map_3D 方法获取坐标映射。
  365.         3.调用 _bilinear_interpolate_3D 方法进行双线性插值，得到变形后的特征。
  366.         4.返回变形后的特征。
  367.         """
  368.         y, x = self._coordinate_map_3D(offset, if_offset)
  369.         deformed_feature = self._bilinear_interpolate_3D(input, y, x)
  370.         return deformed_feature
  372. #---------------------------------YOLOv5 专用部分↓---------------------------------
  373. class DSConv_Bottleneck(nn.Module):
  374.     # DSConv bottleneck
  375.     def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
  376.         super().__init__()
  377.         c_ = int(c2 * e)  # hidden channels
  378.         self.cv1 = Conv(c1, c_, 1, 1)
  379.         self.cv2 = Conv(c_, c2, 3, 1, g=g)
  380.         self.add = shortcut and c1 == c2
  381.         self.snc = DSConv(c2, c2, 3, 1, 1, True)
  383.     def forward(self, x):
  384.         return x + self.snc(self.cv2(self.cv1(x))) if self.add else self.snc(self.cv2(self.cv1(x)))
  387. class DSConv_C3(nn.Module):
  388.     # DSConv Bottleneck with 3 convolutions
  389.     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
  391.         super().__init__()
  392.         c_ = int(c2 * e)  # hidden channels
  393.         self.cv1 = Conv(c1, c_, 1, 1)
  394.         self.cv2 = Conv(c1, c_, 1, 1)
  395.         self.cv3 = Conv(2 * c_, c2, 1)  # act=FReLU(c2)
  396.         self.m = nn.Sequential(*(DSConv_Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)))
  398.     def forward(self, x):
  399.         return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
  401. #---------------------------------YOLOv5 专用部分↑---------------------------------

第②步：修改yolo.py文件  

再来修改yolo.py，在parse_model函数中找到 elif m is nn.BatchNorm2d:语句，在其后面加上下面代码： 

  1. elif m in (DSConv, DSConv_C3):
  2.             c1, c2 = ch[f], args[0]
  3.             if c2 != nc:
  4.                 c2 = make_divisible(c2 * gw, 8)
  5.             args = [c1, c2, *args[1:]]
  6.             if m is DSConv_C3:
  7.                 args.insert(2, n)  # number of repeats
  8.                 n = 1

 如下图所示：

第③步：创建自定义的yaml文件  

第1种，替换conv结构

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  3. # Parameters
  4. nc: 80  # number of classes
  5. depth_multiple: 0.33  # model depth multiple
  6. width_multiple: 0.5  # layer channel multiple
  7. anchors:
  8.   - [10,13, 16,30, 33,23]  # P3/8
  9.   - [30,61, 62,45, 59,119]  # P4/16
  10.   - [116,90, 156,198, 373,326]  # P5/32
  12. # YOLOv5 v6.0 backbone
  13. backbone:
  14.   # [from, number, module, args]
  15.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  16.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  17.    [-1, 3, C3, [128]],
  18.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  19.    [-1, 6, C3, [256]],
  20.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  21.    [-1, 9, C3, [512]],
  22.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  23.    [-1, 3, C3, [1024]],
  24.    [-1, 1, SPPF, [1024, 5]],  # 9
  25.   ]
  27. # YOLOv5 v6.0 head
  28. head:
  29.   [[-1, 1, Conv, [512, 1, 1]],
  30.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  31.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
  32.    [-1, 3, C3, [512]],  # 13
  34.    [-1, 1, Conv, [256, 1, 1]],
  35.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  36.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
  37.    [-1, 3, DSConv, [256, 3,1,1,True]],  # 17 (P3/8-small)
  39.    [-1, 1, Conv, [256, 3, 2]],
  40.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  41.    [-1, 3, DSConv, [512, 3,1,1,True]],  # 20 (P4/16-medium)
  43.    [-1, 1, Conv, [512, 3, 2]],
  44.    [[-1, 10], 1, Concat, [1]],  # cat head P5
  45.    [-1, 3, DSConv, [1024, 3,1,1,True]],  # 23 (P5/32-large)
  47.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  48.   ]

这里要注意一个问题，替换时DSConv参数是需要做对应修改：

如下图栗子所示： 

如果直接改模块名会出现缺参报错：

TypeError: __init__() missing 2 required positional arguments: 'morph' and 'if_offset'

第2种，替换C3模块 

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  3. # Parameters
  4. nc: 80  # number of classes
  5. depth_multiple: 0.33  # model depth multiple
  6. width_multiple: 0.5  # layer channel multiple
  7. anchors:
  8.   - [10,13, 16,30, 33,23]  # P3/8
  9.   - [30,61, 62,45, 59,119]  # P4/16
  10.   - [116,90, 156,198, 373,326]  # P5/32
  12. # YOLOv5 v6.0 backbone
  13. backbone:
  14.   # [from, number, module, args]
  15.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  16.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  17.    [-1, 3, DSConv_C3, [128]],
  18.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  19.    [-1, 6, DSConv_C3, [256]],
  20.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  21.    [-1, 9, DSConv_C3, [512]],
  22.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  23.    [-1, 3, DSConv_C3, [1024]],
  24.    [-1, 1, SPPF, [1024, 5]],  # 9
  25.   ]
  27. # YOLOv5 v6.0 head
  28. head:
  29.   [[-1, 1, Conv, [512, 1, 1]],
  30.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  31.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
  32.    [-1, 3, C3, [512]],  # 13
  34.    [-1, 1, Conv, [256, 1, 1]],
  35.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  36.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
  37.    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
  39.    [-1, 1, Conv, [256, 3, 2]],
  40.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  41.    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
  43.    [-1, 1, Conv, [512, 3, 2]],
  44.    [[-1, 10], 1, Concat, [1]],  # cat head P5
  45.    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
  47.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  48.   ]

替换C3模块直接改模块名字就行。 

第④步：验证是否加入成功

运行yolo.py

第1种 

第2种  

 这样就OK啦！