conv.py
ultralytics\nn\modules\conv.py
目录
conv.py
1.所需的库和模块
2.def autopad(k, p=None, d=1):
3.class Conv(nn.Module):
4.class Conv2(Conv):
5.class LightConv(nn.Module):
6.class DWConv(Conv):
7.class DWConvTranspose2d(nn.ConvTranspose2d):
8.class ConvTranspose(nn.Module):
9.class Focus(nn.Module):
10.class GhostConv(nn.Module):
11.class RepConv(nn.Module):
12.class ChannelAttention(nn.Module):
13.class SpatialAttention(nn.Module):
14.class CBAM(nn.Module):
15.class Concat(nn.Module):
1.所需的库和模块
# Ultralytics YOLO 🚀, AGPL-3.0 license
"""Convolution modules."""
import math
import numpy as np
import torch
import torch.nn as nn
# __all__ = ["function1", "function2", "Class1"]
# __all__ 是一个特殊的变量,用于定义模块的公开接口。当使用 from module import * 这样的导入语句时, __all__ 列表中包含的名称会被导入。如果 __all__ 没有被定义,那么模块中所有以非双下划线( __ )开头的公开名称都会被导入。
# __all__ 变量通常在模块的顶部定义,并且是一个包含字符串的列表或元组。这些字符串是模块中定义的函数、类和变量的名称。定义 __all__ 可以使得模块的接口更加清晰,并且可以控制哪些名称会被 import * 导入。
# 这段代码定义了一个名为 __all__ 的元组,它包含了一系列的类名。这个元组通常用于Python模块中,以明确指出该模块对外公开的接口或类。当某个模块被导入时,如果使用了 from module import * 这样的导入语句, __all__ 元组中列出的名称将会被导入。
__all__ = (
"Conv",
"Conv2",
"LightConv",
"DWConv",
"DWConvTranspose2d",
"ConvTranspose",
"Focus",
"GhostConv",
"ChannelAttention",
"SpatialAttention",
"CBAM",
"Concat",
"RepConv",
)
2.def autopad(k, p=None, d=1):
# 这段代码定义了一个名为 autopad 的函数,它用于自动计算卷积层中所需的填充(padding)值,以确保输出特征图(feature map)的大小与输入特征图相同(即 'same' 填充)。这种填充方式在卷积神经网络中很常见,尤其是在使用卷积层保持特征图尺寸不变时。
# 这一行定义了一个名为 autopad 的函数,它接受三个参数。
# 1.k :代表卷积核的大小。
# 2.p :代表填充的大小,可选。
# 3.d :代表膨胀系数,默认值为1。这个函数的目的是自动计算出为了得到与输入相同大小的输出(same padding)所需要的填充大小。
def autopad(k, p=None, d=1): # kernel, padding, dilation
"""Pad to 'same' shape outputs."""
# 膨胀系数 d 是否大于1。
if d > 1:
# 如果膨胀系数 d 大于1,并且 k 是整数(代表卷积核是正方形的),则计算实际的卷积核大小。如果 k 是列表(代表卷积核可能是非正方形的),则对列表中的每个元素进行相同的计算。这一步是为了考虑到膨胀对卷积核大小的影响。
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
# 检查是否没有提供 p 参数(即填充大小)。
if p is None:
# 如果 p 参数没有提供,函数将自动计算填充大小。如果 k 是整数,填充大小就是卷积核大小的一半;如果 k 是列表,列表中的每个元素都除以2得到对应的填充大小。
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
# 返回计算出的填充大小。
return p
# autopad 函数用于自动计算卷积操作中为了保持输入输出尺寸一致(same padding)所需的填充大小。它考虑了卷积核的大小和膨胀系数,如果用户没有指定填充大小,函数会根据卷积核的大小自动计算。这个函数对于实现卷积神经网络中保持特征图尺寸的操作非常有用。
3.class Conv(nn.Module):
# ✅
# 这段代码定义了一个名为 Conv 的类,它是一个神经网络中的卷积层,继承自PyTorch的 nn.Module 。
# 定义了一个名为 Conv 的类,它继承自PyTorch的 nn.Module ,这是一个基础的模块类,用于构建神经网络。
class Conv(nn.Module):
# 标准卷积,带有参数(ch_in、ch_out、kernel、stride、padding、groups、dilation、activation)。
"""Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
# 定义了一个默认的激活函数,这里使用的是SiLU(也称为Swish),这是一个自门控的激活函数。
default_act = nn.SiLU() # default activation
# 这是 Conv 类的构造函数,它接受多个参数来初始化卷积层。
# 1.c1 和 2.c2 :分别是输入和输出通道数。
# 3.k :是卷积核大小,默认为1。
# 4.s :是步长,默认为1。
# 5.p :是填充大小,可以自动计算,默认为 None。
# 6.g :是组数,默认为1。
# 7.d :是膨胀系数,默认为1。
# 8.act :是激活函数,如果为 True 则使用默认激活函数,如果为 nn.Module 实例则使用该激活函数,否则使用 nn.Identity() (即不使用激活函数)。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
"""Initialize Conv layer with given arguments including activation."""
# 调用父类的构造函数。
super().__init__()
# 创建一个 nn.Conv2d 卷积层实例,使用 autopad 函数自动计算填充大小,不使用偏置项( bias=False )。
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
# 创建一个 nn.BatchNorm2d 批量归一化层实例,用于归一化卷积层的输出。
self.bn = nn.BatchNorm2d(c2)
# 根据 act 参数的值设置激活函数。如果 act 为 True ,则使用默认激活函数;如果 act 是 nn.Module 的实例,则使用该激活函数;否则,使用 nn.Identity() ,即不应用激活函数。
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
# 定义 forward 方法,这是PyTorch模块的前向传播方法,用于定义数据如何通过模块。
def forward(self, x):
# 对输入张量应用卷积、批量标准化和激活。
"""Apply convolution, batch normalization and activation to input tensor."""
# 应用卷积层、批量归一化层和激活函数到输入张量 x ,并返回结果。
return self.act(self.bn(self.conv(x)))
# 定义一个名为 forward_fuse 的方法。
def forward_fuse(self, x):
# 对二维数据执行转置卷积。
"""Perform transposed convolution of 2D data."""
# 这个方法直接将输入张量 x 通过卷积层,然后应用激活函数,并返回结果。注意,这个方法没有使用批量归一化层。
return self.act(self.conv(x))
# Conv 类是一个灵活的卷积层实现,它支持多种参数配置,包括输入输出通道数、卷积核大小、步长、填充、组数、膨胀系数和激活函数。它提供了标准的前向传播方法 forward ,以及一个不使用批量归一化的简化前向传播方法 forward_fuse 。这个类可以作为构建更复杂神经网络结构的基本构件。
4.class Conv2(Conv):
# 这段代码定义了一个名为 Conv2 的新类,它继承自之前定义的 Conv 类,并添加了一些额外的功能。
# 定义了一个名为 Conv2 的类,它继承自 Conv 类。
class Conv2(Conv):
# 带有 Conv 融合的简化 RepConv 模块。
"""Simplified RepConv module with Conv fusing."""
# 这是 Conv2 类的构造函数,它接受与 Conv 类相同的参数,但将卷积核大小 3.k 的默认值改为3。
def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
# 使用给定的参数(包括激活)初始化 Conv 层。
"""Initialize Conv layer with given arguments including activation."""
# 调用父类 Conv 的构造函数,初始化卷积层。
super().__init__(c1, c2, k, s, p, g=g, d=d, act=act)
# 创建一个额外的1x1卷积层 cv2 ,这个卷积层用于在原有的卷积操作后增加一个额外的卷积步骤,可以用于调整通道数或实现残差连接。
self.cv2 = nn.Conv2d(c1, c2, 1, s, autopad(1, p, d), groups=g, dilation=d, bias=False) # add 1x1 conv
# 定义 forward 方法,这是PyTorch模块的前向传播方法。
def forward(self, x):
# 对输入张量应用卷积、批量归一化和激活。
"""Apply convolution, batch normalization and activation to input tensor."""
# 在前向传播中,将原始卷积层 conv 和额外的1x1卷积层 cv2 的输出相加,然后应用批量归一化和激活函数。
return self.act(self.bn(self.conv(x) + self.cv2(x)))
# 定义一个名为 forward_fuse 的方法,用于在不需要额外1x1卷积时的前向传播。
def forward_fuse(self, x):
# 将融合卷积、批量归一化和激活应用于输入张量。
"""Apply fused convolution, batch normalization and activation to input tensor."""
# 在 forward_fuse 方法中,只使用原始卷积层 conv 的输出,并应用批量归一化和激活函数。
return self.act(self.bn(self.conv(x)))
# 定义一个名为 fuse_convs 的方法,用于将两个卷积层融合。
def fuse_convs(self):
# 融合并行卷积。
"""Fuse parallel convolutions."""
# 创建一个与原始卷积层权重相同形状的零张量 w 。
w = torch.zeros_like(self.conv.weight.data)
# 计算 w 的中间索引,用于将1x1卷积层的权重插入到 w 中。
i = [x // 2 for x in w.shape[2:]]
# 将1x1卷积层的权重复制到 w 的中间位置。
w[:, :, i[0] : i[0] + 1, i[1] : i[1] + 1] = self.cv2.weight.data.clone()
# 将修改后的权重 w 加到原始卷积层的权重上,实现两个卷积层的融合。
self.conv.weight.data += w
# object.__delattr__(name)
# 在 Python 中, __delattr__() 是一个魔法方法(magic method),也称为特殊方法或内置方法。它在对象的属性被删除时被自动调用。这个方法有一个默认实现,它简单地调用内置的 delattr() 函数来删除属性。
# 参数 :
# name : 要删除的属性的名称。
# 作用 :
# delattr() 函数用于删除对象的属性。当需要从类实例中删除属性时,Python 会自动调用这个方法。
# 注意事项 :
# 使用 __delattr__() 时要谨慎,因为如果方法内部发生异常,属性可能不会被删除。
# 如果你重写了 __delattr__() 方法,确保最终调用基类的 __delattr__() 方法来确保属性被正确删除,否则可能会出现属性删除失败的情况。
# 删除 cv2 属性,因为它的权重已经被融合到 conv 中。
self.__delattr__("cv2")
# 将 forward 方法替换为 forward_fuse ,使得后续的前向传播不再使用额外的1x1卷积层。
self.forward = self.forward_fuse
# Conv2 类在 Conv 类的基础上添加了一个1x1卷积层,这可以用于实现残差连接或调整通道数。它提供了一个 fuse_convs 方法来融合这两个卷积层,使得在训练完成后可以简化网络结构。这种技术在某些深度学习架构中用于减少模型复杂度和提高效率。
5.class LightConv(nn.Module):
# 这段代码定义了一个名为 LightConv 的类,它继承自PyTorch的 nn.Module 。 LightConv 是一个轻量级的卷积模块,通常用于深度学习模型中以减少参数数量和计算量。
# 定义了一个名为 LightConv 的类,它继承自PyTorch的 nn.Module ,这是一个基础的模块类,用于构建神经网络。
class LightConv(nn.Module):
# 带参数(ch_in、ch_out、kernel)的轻卷积。
"""
Light convolution with args(ch_in, ch_out, kernel).
https://github.com/PaddlePaddle/PaddleDetection/blob/develop/ppdet/modeling/backbones/hgnet_v2.py
"""
# 这是 LightConv 类的构造函数,它接受四个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :深度卷积的卷积核大小,默认为1。
# 4.act :激活函数,默认为ReLU。
def __init__(self, c1, c2, k=1, act=nn.ReLU()):
# 使用给定的参数(包括激活)初始化 Conv 层。
"""Initialize Conv layer with given arguments including activation."""
# 调用父类的构造函数。
super().__init__()
# 创建一个 Conv 实例 conv1 ,这是一个1x1卷积层,用于改变通道数而不改变空间维度。这里 act=False 表示不使用激活函数。
self.conv1 = Conv(c1, c2, 1, act=False)
# 创建一个 DWConv 实例 conv2 ,这是一个深度卷积层(也称为深度可分离卷积),用于在保持参数数量和计算量较低的同时捕获空间信息。
self.conv2 = DWConv(c2, c2, k, act=act)
# 定义 forward 方法,这是PyTorch模块的前向传播方法。
def forward(self, x):
# 对输入张量应用 2 个卷积。
"""Apply 2 convolutions to input tensor."""
# 在前向传播中,首先将输入 x 通过 conv1 ,然后将 conv1 的输出作为 conv2 的输入,并返回 conv2 的输出。
return self.conv2(self.conv1(x))
# LightConv 类是一个轻量级的卷积模块,它由两个卷积层组成 :一个1x1卷积层 conv1 和一个深度卷积层 conv2 。这种结构在需要减少参数数量和计算量时非常有用,同时仍然能够捕获输入数据的重要特征。这种设计在一些高效的深度学习架构中很常见,例如MobileNet等。
6.class DWConv(Conv):
# 这段代码定义了一个名为 DWConv 的类,它继承自之前定义的 Conv 类,并专门用于实现深度可分离卷积(Depthwise Separable Convolution)。
# 定义了一个名为 DWConv 的类,它继承自 Conv 类。
class DWConv(Conv):
# 深度卷积。
"""Depth-wise convolution."""
# 这是 DWConv 类的构造函数,它接受六个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核大小,默认为1。
# 4.s :步长,默认为1。
# 5.d :膨胀系数,默认为1。
# 6.act :激活函数,如果为 True 则使用默认激活函数。
def __init__(self, c1, c2, k=1, s=1, d=1, act=True): # ch_in, ch_out, kernel, stride, dilation, activation
# 使用给定的参数初始化深度卷积。
"""Initialize Depth-wise convolution with given parameters."""
# math.gcd(a, b, *integers)
# math.gcd() 函数是 Python 标准库 math 模块中的一个函数,用于计算两个或多个整数的最大公约数(Greatest Common Divisor,GCD)。最大公约数是能同时整除这些整数的最大正整数。
# 参数说明 :
# a 和 b :两个非负整数,函数将计算它们的最大公约数。
# *integers :额外的整数参数, math.gcd() 可以接受两个以上的整数,并计算所有提供的整数的最大公约数。
# 返回值 :
# 返回 a 和 b (以及任何额外提供的整数)的最大公约数。
# 在深度学习模型中,特别是在实现深度卷积(Depth-wise Convolution)时, math.gcd() 函数可以用来确定卷积层的组数,这样可以使每个输入通道独立地进行卷积操作,从而减少参数数量和计算量。
# 例如,在 DWConv 类的初始化方法中,使用 math.gcd(c1, c2) 来设置组数,其中 c1 和 c2 分别是输入和输出通道的数量。这样做可以确保每个输入通道都与相同的卷积核进行卷积,这是实现深度卷积的关键步骤。
# 调用父类 Conv 的构造函数,初始化深度可分离卷积层。这里 g 参数被设置为 c1 和 c2 的最大公约数,这是深度可分离卷积的一个特点,它将输入通道分组,每组包含 g 个通道,并对每组应用一个卷积核。
super().__init__(c1, c2, k, s, g=math.gcd(c1, c2), d=d, act=act)
# DWConv 类是一个专门用于实现深度可分离卷积的类。深度可分离卷积是一种高效的卷积操作,它将标准的卷积操作分解为两个较小的操作 :深度卷积(每个输入通道独立卷积)和逐点卷积(1x1卷积)。这种分解减少了模型的参数数量和计算量,同时保持了相似的性能。在 DWConv 类中,通过将组数 g 设置为输入和输出通道数的最大公约数,实现了深度卷积的效果。
7.class DWConvTranspose2d(nn.ConvTranspose2d):
# 这段代码定义了一个名为 DWConvTranspose2d 的类,它继承自PyTorch的 nn.ConvTranspose2d ,并专门用于实现深度可分离的转置卷积(Depthwise Separable Transpose Convolution)。
# 定义了一个名为 DWConvTranspose2d 的类,它继承自PyTorch的 nn.ConvTranspose2d ,这是一个基础的转置卷积层类。
class DWConvTranspose2d(nn.ConvTranspose2d):
# 深度转置卷积。
"""Depth-wise transpose convolution."""
# 这是 DWConvTranspose2d 类的构造函数,它接受六个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核大小,默认为1。
# 4.s :步长,默认为1。
# 5.p1 :输入填充大小,默认为0。
# 6.p2 :输出填充大小,默认为0。
def __init__(self, c1, c2, k=1, s=1, p1=0, p2=0): # ch_in, ch_out, kernel, stride, padding, padding_out
# 使用给定的参数初始化 DWConvTranspose2d 类。
"""Initialize DWConvTranspose2d class with given parameters."""
# 调用父类 nn.ConvTranspose2d 的构造函数,初始化深度可分离的转置卷积层。这里 groups 参数被设置为 c1 和 c2 的最大公约数,这是深度可分离卷积的一个特点,它将输入通道分组,每组包含 g 个通道,并对每组应用一个卷积核。
super().__init__(c1, c2, k, s, p1, p2, groups=math.gcd(c1, c2))
# DWConvTranspose2d 类是一个专门用于实现深度可分离的转置卷积的类。转置卷积(也称为反卷积或上采样卷积)是一种用于增加数据空间维度的卷积操作,常用于生成模型或上采样特征图。
# 深度可分离的转置卷积是一种高效的转置卷积操作,它将标准的转置卷积操作分解为两个较小的操作 :深度转置卷积(每个输入通道独立转置卷积)和逐点转置卷积(1x1转置卷积)。这种分解减少了模型的参数数量和计算量,同时保持了相似的性能。
# 在 DWConvTranspose2d 类中,通过将组数 groups 设置为输入和输出通道数的最大公约数,实现了深度转置卷积的效果。
8.class ConvTranspose(nn.Module):
# 这段代码定义了一个名为 ConvTranspose 的类,它继承自PyTorch的 nn.Module ,并用于实现转置卷积层(也称为反卷积层)。
# 定义了一个名为 ConvTranspose 的类,它继承自PyTorch的 nn.Module ,这是一个基础的模块类,用于构建神经网络。
class ConvTranspose(nn.Module):
# 卷积转置 2d 层。
"""Convolution transpose 2d layer."""
# 定义了一个默认的激活函数,这里使用的是SiLU(也称为Swish),这是一个自门控的激活函数。
default_act = nn.SiLU() # default activation
# 这是 ConvTranspose 类的构造函数,它接受七个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核大小,默认为2。
# 4.s :步长,默认为2。
# 5.p :填充大小,默认为0。
# 6.bn :是否使用批量归一化,默认为True。
# 7.act :激活函数,如果为 True 则使用默认激活函数。
def __init__(self, c1, c2, k=2, s=2, p=0, bn=True, act=True):
# 使用批量标准化和激活函数初始化 ConvTranspose2d 层。
"""Initialize ConvTranspose2d layer with batch normalization and activation function."""
# 调用父类的构造函数。
super().__init__()
# 创建一个 nn.ConvTranspose2d 转置卷积层实例,如果 bn 为True,则转置卷积层不使用偏置项( bias=False ),否则使用偏置项( bias=True )。
self.conv_transpose = nn.ConvTranspose2d(c1, c2, k, s, p, bias=not bn)
# 根据 bn 参数的值,创建一个 nn.BatchNorm2d 批量归一化层实例或 nn.Identity (即不使用批量归一化)。
self.bn = nn.BatchNorm2d(c2) if bn else nn.Identity()
# 根据 act 参数的值设置激活函数。如果 act 为 True ,则使用默认激活函数;如果 act 是 nn.Module 的实例,则使用该激活函数;否则,使用 nn.Identity() ,即不应用激活函数。
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
# 定义 forward 方法,这是PyTorch模块的前向传播方法。
def forward(self, x):
# 将转置卷积、批量标准化和激活应用于输入。
"""Applies transposed convolutions, batch normalization and activation to input."""
# 在前向传播中,首先将输入 x 通过转置卷积层 conv_transpose ,然后将输出通过批量归一化层 bn ,最后应用激活函数 act ,并返回结果。
return self.act(self.bn(self.conv_transpose(x)))
# 定义一个名为 forward_fuse 的方法,用于在不需要批量归一化时的前向传播。
def forward_fuse(self, x):
# 将激活和卷积转置操作应用于输入。
"""Applies activation and convolution transpose operation to input."""
# 在 forward_fuse 方法中,直接将输入 x 通过转置卷积层 conv_transpose ,然后应用激活函数 act ,并返回结果。
return self.act(self.conv_transpose(x))
# ConvTranspose 类是一个实现转置卷积操作的类,它提供了批量归一化和激活函数的选项。这个类可以用于实现上采样或转置卷积操作,常用于生成模型或特征图的上采样。通过提供 forward 和 forward_fuse 两个方法, ConvTranspose 类允许用户根据需要选择是否使用批量归一化。
9.class Focus(nn.Module):
# 这段代码定义了一个名为 Focus 的类,它继承自PyTorch的 nn.Module 。 Focus 类的设计意图是实现一种特殊的卷积操作,通常用于图像处理中的特征融合。
# 定义了一个名为 Focus 的类,它继承自PyTorch的 nn.Module ,这是一个基础的模块类,用于构建神经网络。
class Focus(nn.Module):
# 将 wh 信息聚焦到 c 空间。
"""Focus wh information into c-space."""
# 这是 Focus 类的构造函数,它接受七个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核大小,默认为1。
# 4.s :步长,默认为1。
# 5.p :填充大小,默认为 None。
# 6.g :组数,默认为1。
# 7.act :激活函数,如果为 True 则使用默认激活函数。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True):
# 使用用户定义的通道、卷积、填充、组和激活值初始化 Focus 对象。
"""Initializes Focus object with user defined channel, convolution, padding, group and activation values."""
# 调用父类的构造函数。
super().__init__()
# 创建一个 Conv 实例 conv ,这是一个卷积层。这里的输入通道数被设置为 c1 * 4 ,因为 Focus 类的设计意图是将四个子特征图(通过下采样得到的)合并为一个特征图。
self.conv = Conv(c1 * 4, c2, k, s, p, g, act=act)
# 这行代码被注释掉了,但它暗示了 Focus 类可能有一个 contract 属性,用于在空间维度上压缩特征图。 Contract 是一个自定义的模块,用于下采样。
# self.contract = Contract(gain=2)
# 定义 forward 方法,这是PyTorch模块的前向传播方法。
def forward(self, x):
# 将卷积应用于连接张量并返回输出。
# 输入形状为 (b,c,w,h),输出形状为 (b,4c,w/2,h/2)。
"""
Applies convolution to concatenated tensor and returns the output.
Input shape is (b,c,w,h) and output shape is (b,4c,w/2,h/2).
"""
# 在前向传播中, Focus 类首先将输入 x 的四个子特征图(通过不同的索引方式获取)沿着通道维度( dim=1 )拼接起来,然后通过卷积层 conv 进行处理,并返回结果。
return self.conv(torch.cat((x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]), 1))
# 这行代码被注释掉了,但它表明 Focus 类可能有一个备选的前向传播路径,即先通过 contract 模块下采样,然后再通过卷积层 conv 处理。
# return self.conv(self.contract(x))
# Focus 类是一个特征融合模块,它通过将四个子特征图合并为一个特征图来实现特征融合。这种设计在某些卷积神经网络架构中用于整合不同尺度的特征信息。 Focus 类通过 Conv 类实现卷积操作,并且可以通过激活函数对输出进行非线性变换。注释掉的代码表明, Focus 类可能还有其他变体,例如通过 Contract 模块实现下采样。
10.class GhostConv(nn.Module):
# 这段代码定义了一个名为 GhostConv 的类,它继承自PyTorch的 nn.Module 。 GhostConv 是一种用于深度学习中的“幽灵卷积”结构,它通过创建额外的卷积层来增加网络的宽度,而不需要显著增加参数数量。
# 定义了一个名为 GhostConv 的类,它继承自PyTorch的 nn.Module 。
class GhostConv(nn.Module):
# 幽灵卷积。
"""Ghost Convolution https://github.com/huawei-noah/ghostnet."""
# 这是 GhostConv 类的构造函数,它接受六个参数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核大小,默认为1。
# 4.s :步长,默认为1。
# 5.g :组数,默认为1。
# 6.act :激活函数,如果为 True 则使用默认激活函数。
def __init__(self, c1, c2, k=1, s=1, g=1, act=True):
# 使用输入通道、输出通道、内核大小、步幅、组和激活初始化 GhostConv 对象。
"""Initializes the GhostConv object with input channels, output channels, kernel size, stride, groups and
activation.
"""
# 调用父类的构造函数。
super().__init__()
# 计算隐藏通道数 c_ ,这是 c2 的一半。这个值表示 GhostConv 结构中两个卷积层的输出通道数。
c_ = c2 // 2 # hidden channels
# 创建第一个 Conv 实例 cv1 ,这是一个标准卷积层,用于将输入通道数 c1 转换为隐藏通道数 c_ 。
self.cv1 = Conv(c1, c_, k, s, None, g, act=act)
# 创建第二个 Conv 实例 cv2 ,这是一个标准卷积层,用于将隐藏通道数 c_ 再次转换为隐藏通道数 c_ 。这里使用了5x5的卷积核和1的步长。
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act=act)
# 定义 forward 方法,这是PyTorch模块的前向传播方法。
def forward(self, x):
# 通过带有跳过连接的 Ghost Bottleneck 层进行前向传播。
"""Forward propagation through a Ghost Bottleneck layer with skip connection."""
# 在前向传播中,首先将输入 x 通过第一个卷积层 cv1 。
y = self.cv1(x)
# 然后,将 cv1 的输出 y 和 cv2 的输出沿着通道维度( dim=1 )拼接起来,并返回结果。
return torch.cat((y, self.cv2(y)), 1)
# GhostConv 类实现了一个“幽灵卷积”结构,它通过两个卷积层来增加网络的宽度,而不需要显著增加参数数量。这种结构的第一个卷积层将输入特征图的通道数减少到一半,第二个卷积层保持通道数不变。最后,两个卷积层的输出沿着通道维度拼接起来,从而实现了网络宽度的增加。这种设计可以在不显著增加计算负担的情况下提高网络的表达能力。
11.class RepConv(nn.Module):
# 这段代码定义了一个名为 RepConv 的类,它继承自PyTorch的 nn.Module 。 RepConv 是一种可重构的卷积层,它允许在训练时使用两个卷积层(一个3x3卷积和一个1x1卷积),而在部署时将它们合并为一个单一的卷积层以减少计算量和模型大小。
# 定义了一个名为 RepConv 的类,它继承自PyTorch的 nn.Module 。
class RepConv(nn.Module):
# RepConv 是一个基本的 rep 样式块,包括训练和部署状态。
# 此模块用于 RT-DETR。
"""
RepConv is a basic rep-style block, including training and deploy status.
This module is used in RT-DETR.
Based on https://github.com/DingXiaoH/RepVGG/blob/main/repvgg.py
"""
# 定义了一个默认的激活函数,这里使用的是SiLU(也称为Swish),这是一个自门控的激活函数。
default_act = nn.SiLU() # default activation
# 这段代码是 RepConv 的构造函数 __init__ 的定义。这个构造函数初始化了一个可重构的卷积层,允许在训练时使用两个卷积层(一个3x3卷积和一个1x1卷积),而在部署时将它们合并为一个单一的卷积层。
# 这是构造函数的定义,它接受以下参数 :
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.k : 卷积核大小,默认为3。
# 4.s : 步长,默认为1。
# 5.p : 填充大小,默认为1。
# 6.g : 组数,默认为1。
# 7.d : 膨胀系数,默认为1(尽管在这个构造函数中没有使用)。
# 8.act : 是否使用激活函数,默认为True。
# 9.bn : 是否使用批量归一化,默认为False。
# 10.deploy : 是否为部署模式,默认为False。
def __init__(self, c1, c2, k=3, s=1, p=1, g=1, d=1, act=True, bn=False, deploy=False):
# 使用输入、输出和可选激活函数初始化光卷积层。
"""Initializes Light Convolution layer with inputs, outputs & optional activation function."""
# 调用父类( nn.Module )的构造函数。
super().__init__()
# 断言确保卷积核大小 k 为3,填充 p 为1,这是RepConv设计的前提条件。
assert k == 3 and p == 1
# 保存 组数 g , 输入通道数 c1 和 输出通道数 c2 为实例变量。
self.g = g
self.c1 = c1
self.c2 = c2
# 根据 act 参数的值设置激活函数。如果 act 为True,则使用默认激活函数;如果 act 是 nn.Module 的实例,则使用该激活函数;否则,使用 nn.Identity() ,即不应用激活函数。
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
# 如果使用批量归一化( bn=True )且输出通道数等于输入通道数( c2 == c1 )且步长为1( s == 1 ),则创建一个批量归一化层 nn.BatchNorm2d 。
self.bn = nn.BatchNorm2d(num_features=c1) if bn and c2 == c1 and s == 1 else None
# 创建第一个 Conv 实例 conv1 ,这是一个3x3卷积层,不包含激活函数。
self.conv1 = Conv(c1, c2, k, s, p=p, g=g, act=False)
# 创建第二个 Conv 实例 conv2 ,这是一个1x1卷积层,不包含激活函数。填充 p 被调整为 p - k // 2 以保持输出尺寸的一致性。
self.conv2 = Conv(c1, c2, 1, s, p=(p - k // 2), g=g, act=False)
# 这个构造函数初始化了 RepConv 类的主要组件,包括两个卷积层和一个可选的批量归一化层。这些组件将被用于前向传播,以及在部署模式下将两个卷积层合并为一个单一的卷积层。
# 这段代码定义了 RepConv 类中的 forward_fuse 方法,该方法用于在部署模式下执行前向传播。
# forward_fuse 方法接受一个输入 1.x ,这个方法的目的是简化模型的前向传播过程,通常用于模型部署阶段。
def forward_fuse(self, x):
"""Forward process."""
# 输入 x 首先通过一个卷积层 self.conv ,然后通过激活函数 self.act 。这里 self.conv 是合并后的卷积层,它包含了原始3x3卷积和1x1卷积的效果,而 self.act 是激活函数。最终,方法返回经过卷积和激活函数处理后的输出。
return self.act(self.conv(x))
# self.conv 是在 fuse_convs 方法中创建的,该方法将 self.conv1 和 self.conv2 (以及可能的批量归一化层)合并为一个单一的卷积层 self.conv 。
# 这样, forward_fuse 方法可以使用这个合并后的卷积层来简化模型的前向传播过程,减少模型的复杂度和计算量,同时保持模型的性能。这种方法特别适用于模型部署,因为它可以减少模型大小,提高推理速度。
# 这段代码定义了 RepConv 类中的 forward 方法,该方法用于在训练模式下执行前向传播。
# forward 方法接受一个输入 1.x ,这个方法的目的是处理输入数据,并通过 RepConv 层的两个分支(3x3卷积和1x1卷积)进行前向传播。
def forward(self, x):
"""Forward process."""
# 首先检查是否存在批量归一化层 self.bn 。如果存在,将输入 x 通过批量归一化层处理,否则 id_out 为0。这里的 id_out 代表批量归一化层的输出,它将被用作残差连接的一部分。
id_out = 0 if self.bn is None else self.bn(x)
# 然后,输入 x 分别通过两个卷积层 self.conv1 (3x3卷积)和 self.conv2 (1x1卷积)。这两个卷积层的输出相加,再加上批量归一化层的输出(如果有),形成残差连接。最后,将这个总和通过激活函数 self.act 处理,返回最终的输出。
return self.act(self.conv1(x) + self.conv2(x) + id_out)
# 这里的设计思想是利用残差连接来提高训练的稳定性,特别是在深度网络中。残差连接允许网络学习到恒等映射,这有助于解决梯度消失的问题。在 RepConv 中,1x1卷积用于增加或减少通道数,而3x3卷积用于提取特征,批量归一化层(如果存在)用于规范化特征。通过这种方式, RepConv 能够在保持性能的同时减少参数数量和计算量。
# 这段代码定义了 RepConv 类中的 get_equivalent_kernel_bias 方法,该方法用于获取两个卷积层(3x3卷积和1x1卷积)以及可能的批量归一化层的等效卷积核和偏置。
# get_equivalent_kernel_bias 方法用于获取合并后的卷积核和偏置。
def get_equivalent_kernel_bias(self):
# 通过将 3x3 核、1x1 核和恒等核及其偏差相加,返回等效核和偏差。
"""Returns equivalent kernel and bias by adding 3x3 kernel, 1x1 kernel and identity kernel with their biases."""
# 调用私有方法 _fuse_bn_tensor 来处理第一个卷积层 self.conv1 (3x3卷积),获取其等效的卷积核 kernel3x3 和偏置 bias3x3 。这个方法会考虑卷积层和其后的批量归一化层(如果存在)。
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.conv1)
# 同样地,处理第二个卷积层 self.conv2 (1x1卷积),获取其等效的卷积核 kernel1x1 和偏置 bias1x1 。
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.conv2)
# 如果存在批量归一化层 self.bn ,则处理它以获取等效的卷积核 kernelid 和偏置 biasid 。这里的 kernelid 实际上是一个恒等映射,用于表示批量归一化层的效果。
kernelid, biasid = self._fuse_bn_tensor(self.bn)
# 最后,将三个等效卷积核相加( kernel3x3 、 kernel1x1 扩展为3x3大小后、 kernelid ),以及将三个偏置相加,得到最终的等效卷积核和偏置,并返回它们。
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
# 这里的关键是将两个不同大小的卷积核(3x3和1x1)以及可能的批量归一化层合并为一个单一的3x3卷积核。这样做的目的是为了在模型部署时减少模型复杂度,通过一个单一的卷积操作替代原来的两个卷积操作,从而提高推理速度。 self._pad_1x1_to_3x3_tensor 方法用于将1x1卷积核扩展为3x3大小,以便于与3x3卷积核相加。
# 这段代码定义了 RepConv 类中的一个私有方法 _pad_1x1_to_3x3_tensor ,该方法用于将1x1卷积核扩展为3x3卷积核。
# _pad_1x1_to_3x3_tensor 方法接受一个参数。
# 1.kernel1x1 :这是1x1卷积层的卷积核。
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
# 将 1x1 张量填充为 3x3 张量。
"""Pads a 1x1 tensor to a 3x3 tensor."""
# 如果 kernel1x1 为 None ,意味着没有1x1卷积核,因此方法返回0。
if kernel1x1 is None:
return 0
# 如果 kernel1x1 不是 None ,则使用 torch.nn.functional.pad 函数将1x1卷积核在四个方向(上下左右)各填充1个单位,从而将其扩展为3x3卷积核。填充列表 [1, 1, 1, 1] 分别对应于(左、右、上、下)的填充量。
else:
# torch.nn.functional.pad(input, pad, mode='constant', value=0)
# torch.nn.functional.pad 是 PyTorch 中的一个函数,用于对输入的张量(tensor)进行填充(padding)。这个函数非常灵活,允许你在不同的维度上应用不同大小的填充。
# 参数说明 :
# input ( torch.Tensor ) : 需要被填充的输入张量。
# pad ( int 或 tuple ) : 指定填充的大小。
# 如果是 int ,表示在所有边界上应用相同的填充大小。
# 如果是 tuple ,应该遵循 (左, 右, 上, 下) 的顺序,对于多维数据,可以扩展为 (左, 右, 上, 下, 前, 后) 。
# mode ( str , 可选) : 指定填充的模式。
# 'constant' : 常数填充,默认值,使用 value 参数指定的值进行填充。
# 'reflect' : 反射填充,类似于图像处理中的反射效果。
# 'replicate' : 复制边缘值进行填充。
# value ( float 或 int , 可选) : 当 mode 为 'constant' 时,用于指定填充的常数值,默认为 0。
# 返回值 :
# 返回一个新的张量,该张量是输入张量经过指定填充后的结果。
# torch.nn.functional.pad 函数在深度学习中非常有用,尤其是在需要对输入数据进行预处理或在自定义层中调整张量尺寸时。
return torch.nn.functional.pad(kernel1x1, [1, 1, 1, 1])
# 这个方法是 get_equivalent_kernel_bias 中的一个步骤,用于将1x1卷积核调整为与3x3卷积核相同的尺寸,以便可以将它们相加,从而创建一个等效的3x3卷积核。这种技术允许我们将两个不同尺寸的卷积核合并为一个,这对于模型简化和部署非常有用。
# 这段代码定义了 RepConv 类中的一个私有方法 _fuse_bn_tensor ,该方法用于将卷积层( Conv )或批量归一化层( nn.BatchNorm2d )的权重和偏置融合到一起,以便可以与卷积核合并。
# _fuse_bn_tensor 方法接受一个参数。
# 1.branch :这可以是一个卷积层或批量归一化层。
def _fuse_bn_tensor(self, branch):
# 通过融合神经网络的分支来生成适合卷积的内核和偏差。
"""Generates appropriate kernels and biases for convolution by fusing branches of the neural network."""
# 如果 branch 为 None ,则返回两个0值,表示没有卷积核和偏置。
if branch is None:
return 0, 0
# 检查 branch 是否是 Conv 类的实例。
if isinstance(branch, Conv):
# 如果是 Conv 实例,提取卷积层的权重 kernel 和批量归一化层的参数 : running_mean 、 running_var 、 gamma 、 beta 和 eps 。
# 提取了 Conv 实例中卷积层的 权重 。 branch.conv.weight 是一个包含了卷积核参数的PyTorch张量,这些参数是学习得到的,用于在卷积操作中应用。
kernel = branch.conv.weight
# 提取了 Conv 实例中批量归一化层的 运行均值 。 running_mean 是一个张量,存储了在训练过程中计算的均值,用于归一化。
running_mean = branch.bn.running_mean
# 提取了 Conv 实例中批量归一化层的 运行方差 。 running_var 是一个张量,存储了在训练过程中计算的方差,用于归一化。
running_var = branch.bn.running_var
# 提取了 Conv 实例中批量归一化层的 gamma 参数。 gamma 是一个张量,也称为缩放因子,用于在归一化后缩放特征。
gamma = branch.bn.weight
# 提取了 Conv 实例中批量归一化层的 beta 参数。 beta 是一个张量,也称为偏移因子,用于在归一化和缩放后添加偏移。
beta = branch.bn.bias
# 提取了 Conv 实例中批量归一化层的 eps 参数。 eps (epsilon)是一个小的常数,用于防止除以零的情况,并提供数值稳定性。
eps = branch.bn.eps
# 检查 branch 是否是 nn.BatchNorm2d 类的实例。
elif isinstance(branch, nn.BatchNorm2d):
# 如果类中没有 id_tensor 属性,这意味着我们需要创建一个表示恒等映射的卷积核。
if not hasattr(self, "id_tensor"):
# 创建一个3x3的零矩阵,并在对角线上设置为1,以表示恒等映射。然后将这个矩阵转换为PyTorch张量,并移动到批量归一化层权重所在的设备。
# 计算每个组的 输入维度 ( input_dim )。由于 self.c1 是输入通道数, self.g 是组数,所以 input_dim 表示每个组的通道数。
input_dim = self.c1 // self.g
# 使用NumPy创建一个形状为 (self.c1, input_dim, 3, 3) 的零矩阵,其中 self.c1 是输入通道数, input_dim 是每个组的通道数,3x3是卷积核的大小。 dtype=np.float32 指定了矩阵中元素的数据类型为32位浮点数。
kernel_value = np.zeros((self.c1, input_dim, 3, 3), dtype=np.float32)
# 这个循环遍历输入通道,并将每个通道对应的3x3卷积核中心位置设置为1。由于卷积核是3x3的,所以中心位置是(1, 1)。这样,每个通道的卷积核在中心有一个1,其他地方都是0,表示该通道的输入将直接传递到输出,而不会与其它通道混合。
for i in range(self.c1):
kernel_value[i, i % input_dim, 1, 1] = 1
# 将NumPy数组 kernel_value 转换为PyTorch张量,并将其移动到与 branch.weight 相同的设备上(例如CPU或GPU)。这个张量 self.id_tensor 表示恒等映射的卷积核,它将用于表示批量归一化层的效果,或者在没有卷积层时表示恒等操作。
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
# 对于批量归一化层,使用恒等映射 kernel ,并提取批量归一化层的参数。
# 将 self.id_tensor 赋值给 kernel 变量。 self.id_tensor 是一个表示恒等映射的卷积核,它的形状与 branch 的权重形状相匹配,并且用于模拟批量归一化层的效果。
kernel = self.id_tensor
# 提取了批量归一化层的运行均值( running_mean )。这是一个张量,存储了在训练过程中计算的均值,用于归一化。
running_mean = branch.running_mean
# 提取了批量归一化层的运行方差( running_var )。这是一个张量,存储了在训练过程中计算的方差,用于归一化。
running_var = branch.running_var
# 提取了批量归一化层的 gamma 参数,也称为缩放因子。 gamma 是一个张量,用于在归一化后缩放特征。
gamma = branch.weight
# 提取了批量归一化层的 beta 参数,也称为偏移因子。 beta 是一个张量,用于在归一化和缩放后添加偏移。
beta = branch.bias
# 提取了批量归一化层的 eps 参数。 eps (epsilon)是一个小的常数,用于防止除以零的情况,并提供数值稳定性。
eps = branch.eps
# 计算批量归一化层的标准差 std 。
std = (running_var + eps).sqrt()
# 计算缩放因子 t ,并将 gamma 除以 std ,然后重塑为卷积核的形状。
t = (gamma / std).reshape(-1, 1, 1, 1)
# 返回融合后的卷积核( kernel * t )和偏置( beta - running_mean * gamma / std )。
return kernel * t, beta - running_mean * gamma / std
# 这个方法的目的是将卷积层和批量归一化层的参数融合,以便可以将它们视为一个单一的卷积操作。这种融合对于将多个卷积层合并为一个卷积层特别有用,这可以简化模型并减少参数数量。
# 这段代码定义了 RepConv 类中的 fuse_convs 方法,该方法用于将 RepConv 实例中的两个卷积层( conv1 和 conv2 )以及可能的批量归一化层( bn )合并为一个单一的卷积层。
# fuse_convs 方法用于合并卷积层。
def fuse_convs(self):
# 将两个卷积层合并为一个层,并从类中删除未使用的属性。
"""Combines two convolution layers into a single layer and removes unused attributes from the class."""
# 如果 RepConv 实例已经有了一个合并后的卷积层( conv ),则直接返回,不再进行合并操作。
if hasattr(self, "conv"):
return
# 调用 get_equivalent_kernel_bias 方法来获取等效的卷积核和偏置,这些将用于创建新的合并后的卷积层。
kernel, bias = self.get_equivalent_kernel_bias()
# 创建一个新的 nn.Conv2d 卷积层实例。
# 设置新卷积层的参数,包括输入通道数、输出通道数、卷积核大小、步长、填充、膨胀和组数。这些参数都从 conv1 层获取。同时,设置 bias=True 以包含偏置项,并设置 requires_grad_(False) 以禁用梯度计算。
self.conv = nn.Conv2d(
in_channels=self.conv1.conv.in_channels,
out_channels=self.conv1.conv.out_channels,
kernel_size=self.conv1.conv.kernel_size,
stride=self.conv1.conv.stride,
padding=self.conv1.conv.padding,
dilation=self.conv1.conv.dilation,
groups=self.conv1.conv.groups,
bias=True,
).requires_grad_(False)
# 将等效的 卷积核 和 偏置 赋值给新创建的卷积层。
self.conv.weight.data = kernel
self.conv.bias.data = bias
# 遍历 RepConv 实例的所有参数,并将它们从计算图中分离,这样它们就不会参与梯度计算。
for para in self.parameters():
para.detach_()
# 删除原来的两个卷积层 conv1 和 conv2 。
self.__delattr__("conv1")
self.__delattr__("conv2")
# 如果存在名为 nm 的属性,也将其删除。
if hasattr(self, "nm"):
self.__delattr__("nm")
# 如果存在批量归一化层 bn ,也将其删除。
if hasattr(self, "bn"):
self.__delattr__("bn")
# 如果存在恒等映射张量 id_tensor ,也将其删除。
if hasattr(self, "id_tensor"):
self.__delattr__("id_tensor")
# fuse_convs 方法的目的是简化模型结构,将多个卷积和归一化层合并为一个单一的卷积层,这样可以减少模型的参数数量和计算量,同时保持模型的性能。这种合并操作通常在模型训练完成后进行,用于模型的部署和推理。
# RepConv 类是一个可重构的卷积层,它在训练时使用两个卷积层(一个3x3卷积和一个1x1卷积),而在部署时将它们合并为一个单一的卷积层。这种设计可以在训练时提供更好的灵活性和性能,在部署时减少计算量和模型大小。 RepConv 类提供了 fuse_convs 方法来合并卷积层,并提供了 forward 和 forward_fuse 方法来处理前向传播。
12.class ChannelAttention(nn.Module):
# 这段代码定义了一个名为 ChannelAttention 的类,它继承自PyTorch的 nn.Module 。 ChannelAttention 是一个通道注意力模块,通常用于深度学习中的图像处理任务,以增强特定通道的特征并抑制不重要的通道特征。
# 定义了一个名为 ChannelAttention 的类,它继承自PyTorch的 nn.Module 。
class ChannelAttention(nn.Module):
# 通道注意力模块 https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet。
"""Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet."""
# 这是 ChannelAttention 类的构造函数,它接受一个参数。
# 1.channels :表示输入特征图的通道数。
def __init__(self, channels: int) -> None:
# 初始化类并设置所需的基本配置和实例变量。
"""Initializes the class and sets the basic configurations and instance variables required."""
# 调用父类的构造函数。
super().__init__()
# 创建一个自适应平均池化层 self.pool ,它将输入特征图的每个通道的空间维度(高和宽)减少到1x1。
self.pool = nn.AdaptiveAvgPool2d(1)
# 创建一个一维卷积层 self.fc ,它将输入特征图的通道数从 channels 减少到 channels ,实质上是一个全连接层,用于学习通道间的权重。
self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
# 定义一个激活函数 self.act ,这里使用的是Sigmoid函数,它将输出值压缩到0和1之间,用于控制通道的注意力权重。
self.act = nn.Sigmoid()
# 定义 forward 方法,这是PyTorch模块的前向传播方法,接受一个输入张量 x 。
def forward(self, x: torch.Tensor) -> torch.Tensor:
# 使用激活对输入的卷积应用前向传递,可选择使用批量标准化。
"""Applies forward pass using activation on convolutions of the input, optionally using batch normalization."""
# 在前向传播中,首先将输入 x 通过自适应平均池化层 self.pool ,然后通过一维卷积层 self.fc ,接着通过Sigmoid激活函数 self.act 得到通道注意力权重,最后将这些权重乘以原始输入 x ,实现通道注意力的效果。
return x * self.act(self.fc(self.pool(x)))
# ChannelAttention 类实现了通道注意力机制,它通过自适应平均池化和一维卷积学习通道间的权重,然后使用Sigmoid激活函数将这些权重应用到原始输入特征图上,从而增强重要通道的特征并抑制不重要的通道特征。这种机制有助于提高模型对关键信息的敏感度,常用于各种图像识别和分割任务中。
13.class SpatialAttention(nn.Module):
# 这段代码定义了一个名为 SpatialAttention 的类,它继承自PyTorch的 nn.Module 。 SpatialAttention 是一个空间注意力模块,用于强调图像中的重要空间位置并抑制不重要的位置。
# 定义了一个名为 SpatialAttention 的类,它继承自PyTorch的 nn.Module 。
class SpatialAttention(nn.Module):
# 空间注意力模块。
"""Spatial-attention module."""
# 这是 SpatialAttention 类的构造函数,它接受一个参数。
# 1.kernel_size :表示卷积核的大小,默认值为7。
def __init__(self, kernel_size=7):
# 使用内核大小参数初始化空间注意模块。
"""Initialize Spatial-attention module with kernel size argument."""
# 调用父类的构造函数。
super().__init__()
# 断言确保 kernel_size 只能是3或7,这是该模块设计的限制。
assert kernel_size in (3, 7), "kernel size must be 3 or 7" # 卷积核大小必须为 3 或 7 。
# 根据 kernel_size 计算填充大小 padding ,以保持输出尺寸与输入尺寸相同。
padding = 3 if kernel_size == 7 else 1
# 创建一个二维卷积层 self.cv1 ,它接受两个通道的输入(空间特征),输出一个通道(空间注意力权重),卷积核大小为 kernel_size ,填充为 padding ,不包含偏置项。
self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
# 定义一个激活函数 self.act ,这里使用的是Sigmoid函数,它将输出值压缩到0和1之间,用于控制空间位置的注意力权重。
self.act = nn.Sigmoid()
# 定义 forward 方法,这是PyTorch模块的前向传播方法,接受一个输入张量 x 。
def forward(self, x):
# 对输入应用通道和空间注意力进行特征重新校准。
"""Apply channel and spatial attention on input for feature recalibration."""
# 在前向传播中,首先计算输入 x 的均值和最大值,这两个操作分别提取了全局和局部的空间特征。然后,将这两个特征沿着通道维度( dim=1 )拼接起来,并通过卷积层 self.cv1 和Sigmoid激活函数 self.act 得到空间注意力权重,最后将这些权重乘以原始输入 x ,实现空间注意力的效果。
return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))
# SpatialAttention 类实现了空间注意力机制,它通过计算输入特征图的均值和最大值来捕获全局和局部的空间特征,然后使用卷积层和Sigmoid激活函数学习空间位置的权重,最后将这些权重应用到原始输入特征图上,从而强调图像中的重要空间位置并抑制不重要的位置。这种机制有助于提高模型对图像中关键区域的敏感度,常用于各种图像识别和分割任务中。
14.class CBAM(nn.Module):
# 这段代码定义了一个名为 CBAM 的类,它继承自PyTorch的 nn.Module 。 CBAM 代表Convolutional Block Attention Module,是一个集成了通道注意力和空间注意力的模块,用于增强卷积神经网络中的特征表示。
# 定义了一个名为 CBAM 的类,它继承自PyTorch的 nn.Module 。
class CBAM(nn.Module):
# 卷积块注意力模块。
"""Convolutional Block Attention Module."""
# 这是 CBAM 类的构造函数,它接受两个参数。
# 1.c1 :表示输入特征图的通道数。
# 2.kernel_size :表示空间注意力模块中卷积核的大小,默认值为7。
def __init__(self, c1, kernel_size=7):
# 使用给定的输入通道(c1)和内核大小初始化 CBAM。
"""Initialize CBAM with given input channel (c1) and kernel size."""
# 调用父类的构造函数。
super().__init__()
# 创建一个 ChannelAttention 实例,用于实现通道注意力机制。
self.channel_attention = ChannelAttention(c1)
# 创建一个 SpatialAttention 实例,用于实现空间注意力机制,卷积核大小由 kernel_size 参数指定。
self.spatial_attention = SpatialAttention(kernel_size)
# 定义 forward 方法,这是PyTorch模块的前向传播方法,接受一个输入张量 x 。
def forward(self, x):
# 通过 C1 模块应用前向传递。
"""Applies the forward pass through C1 module."""
# 在前向传播中,首先将输入 x 通过通道注意力模块 self.channel_attention ,然后再通过空间注意力模块 self.spatial_attention ,最终返回经过注意力增强的特征图。
return self.spatial_attention(self.channel_attention(x))
# CBAM 类结合了通道注意力和空间注意力,首先通过 ChannelAttention 模块对特征图的通道进行加权,强调重要的通道并抑制不重要的通道;接着,通过 SpatialAttention 模块对特征图的空间维度进行加权,强调重要的空间位置并抑制不重要的位置。这种双重注意力机制有助于提升模型对关键信息的捕捉能力,广泛应用于图像分类、目标检测等任务中,以提高模型的性能和泛化能力。
15.class Concat(nn.Module):
# ✅
# 这段代码定义了一个名为 Concat 的类,它继承自PyTorch的 nn.Module 。 Concat 类用于在指定的维度上对输入的多个张量进行拼接。
# 定义了一个名为 Concat 的类,它继承自PyTorch的 nn.Module 。
class Concat(nn.Module):
# 沿维度连接张量列表。
"""Concatenate a list of tensors along dimension."""
# 这是 Concat 类的构造函数,它接受一个参数。
# 1.dimension :表示张量拼接的维度,默认值为1。
def __init__(self, dimension=1):
# 沿指定维度连接张量列表。
"""Concatenates a list of tensors along a specified dimension."""
# 调用父类的构造函数。
super().__init__()
# 将传入的 dimension 参数保存为实例变量 self.d ,用于后续的张量拼接操作。
self.d = dimension
# 定义 forward 方法,这是PyTorch模块的前向传播方法,接受一个或多个输入张量 x 。
def forward(self, x):
# YOLOv8 mask Proto 模块的前向传递。
"""Forward pass for the YOLOv8 mask Proto module."""
# 在前向传播中,使用PyTorch的 torch.cat 函数将输入的张量 x 在 self.d 指定的维度上进行拼接,并返回结果。
return torch.cat(x, self.d)
# Concat 类提供了一个简单的方法来在指定维度上拼接多个张量。这个类在构建神经网络时非常有用,尤其是在需要将不同来源的特征图或数据合并到一起时。通过指定 dimension 参数,用户可以控制张量在哪个维度上进行拼接,从而实现灵活的特征融合。