YOLOv6-4.0部分代码阅读笔记-efficientrep.py
efficientrep.py
yolov6\models\efficientrep.py
目录
efficientrep.py
1.所需的库和模块
2.class EfficientRep(nn.Module):
3.class EfficientRep6(nn.Module):
4.class CSPBepBackbone(nn.Module):
5.class CSPBepBackbone_P6(nn.Module):
6.class Lite_EffiBackbone(nn.Module):
1.所需的库和模块
from pickle import FALSE
from torch import nn
from yolov6.layers.common import BottleRep, RepVGGBlock, RepBlock, BepC3, SimSPPF, SPPF, SimCSPSPPF, CSPSPPF, ConvBNSiLU, \
MBLABlock, ConvBNHS, Lite_EffiBlockS2, Lite_EffiBlockS12.class EfficientRep(nn.Module):
class EfficientRep(nn.Module):
# EfficientRep 骨干
# EfficientRep 是由硬件感知神经网络设计手工制作的。
# EfficientRep 采用 rep-style 结构，对高计算硬件（例如 GPU）非常友好。
'''EfficientRep Backbone
EfficientRep is handcrafted by hardware-aware neural network design.
With rep-style struct, EfficientRep is friendly to high-computation hardware(e.g. GPU).
'''
# 参数 ：
# 1.in_channels ：输入数据的通道数，默认为3（例如RGB图像）。
# 2.channels_list ：一个列表，定义了每个阶段的输出通道数。
# 3.num_repeats ：一个列表，定义了每个阶段重复块的次数。
# 4.block ：用于构建网络的基本块，默认为 RepVGGBlock 。
# 5.fuse_P2 ：一个布尔值，指示是否融合第二阶段的特征。
# 6.cspsppf ：一个布尔值，指示是否使用 CSPSPPF（Cross Stage Partial Spatial Pyramid Pooling and Fusion）。
def __init__(
self,
in_channels=3,
channels_list=None,
num_repeats=None,
block=RepVGGBlock,
fuse_P2=False,
cspsppf=False
):
# 调用 super().__init__() 初始化父类 nn.Module 。
super().__init__()
# 断言 channels_list 和 num_repeats 不为 None 。
assert channels_list is not None
assert num_repeats is not None
# 初始化 self.fuse_P2 属性。
self.fuse_P2 = fuse_P2
self.stem = block(
in_channels=in_channels,
out_channels=channels_list[0],
kernel_size=3,
stride=2
)
self.ERBlock_2 = nn.Sequential(
block(
in_channels=channels_list[0],
out_channels=channels_list[1],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[1],
out_channels=channels_list[1],
n=num_repeats[1],
block=block,
)
)
self.ERBlock_3 = nn.Sequential(
block(
in_channels=channels_list[1],
out_channels=channels_list[2],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[2],
out_channels=channels_list[2],
n=num_repeats[2],
block=block,
)
)
self.ERBlock_4 = nn.Sequential(
block(
in_channels=channels_list[2],
out_channels=channels_list[3],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[3],
out_channels=channels_list[3],
n=num_repeats[3],
block=block,
)
)
channel_merge_layer = SPPF if block == ConvBNSiLU else SimSPPF
if cspsppf:
channel_merge_layer = CSPSPPF if block == ConvBNSiLU else SimCSPSPPF
self.ERBlock_5 = nn.Sequential(
block(
in_channels=channels_list[3],
out_channels=channels_list[4],
kernel_size=3,
stride=2,
),
RepBlock(
in_channels=channels_list[4],
out_channels=channels_list[4],
n=num_repeats[4],
block=block,
),
channel_merge_layer(
in_channels=channels_list[4],
out_channels=channels_list[4],
kernel_size=5
)
)
# 定义前向传播函数。
def forward(self, x):
# 通过每个阶段的网络，并将特征图添加到输出列表 outputs 。
outputs = []
x = self.stem(x)
x = self.ERBlock_2(x)
# 如果 self.fuse_P2 为 True ，则将第二阶段的特征图添加到输出列表。
if self.fuse_P2:
outputs.append(x)
x = self.ERBlock_3(x)
outputs.append(x)
x = self.ERBlock_4(x)
outputs.append(x)
x = self.ERBlock_5(x)
outputs.append(x)
# 返回一个包含所有输出特征图的元组。
# forward 方法返回一个元组，包含网络各阶段的输出特征图。
return tuple(outputs)3.class EfficientRep6(nn.Module):
class EfficientRep6(nn.Module):
# EfficientRep+P6 主干
# EfficientRep 是由硬件感知神经网络设计手工制作的。
# EfficientRep 采用 rep-style 结构，对高计算硬件（例如 GPU）非常友好。
'''EfficientRep+P6 Backbone
EfficientRep is handcrafted by hardware-aware neural network design.
With rep-style struct, EfficientRep is friendly to high-computation hardware(e.g. GPU).
'''
def __init__(
self,
in_channels=3,
channels_list=None,
num_repeats=None,
block=RepVGGBlock,
fuse_P2=False,
cspsppf=False
):
super().__init__()
assert channels_list is not None
assert num_repeats is not None
self.fuse_P2 = fuse_P2
self.stem = block(
in_channels=in_channels,
out_channels=channels_list[0],
kernel_size=3,
stride=2
)
self.ERBlock_2 = nn.Sequential(
block(
in_channels=channels_list[0],
out_channels=channels_list[1],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[1],
out_channels=channels_list[1],
n=num_repeats[1],
block=block,
)
)
self.ERBlock_3 = nn.Sequential(
block(
in_channels=channels_list[1],
out_channels=channels_list[2],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[2],
out_channels=channels_list[2],
n=num_repeats[2],
block=block,
)
)
self.ERBlock_4 = nn.Sequential(
block(
in_channels=channels_list[2],
out_channels=channels_list[3],
kernel_size=3,
stride=2
),
RepBlock(
in_channels=channels_list[3],
out_channels=channels_list[3],
n=num_repeats[3],
block=block,
)
)
self.ERBlock_5 = nn.Sequential(
block(
in_channels=channels_list[3],
out_channels=channels_list[4],
kernel_size=3,
stride=2,
),
RepBlock(
in_channels=channels_list[4],
out_channels=channels_list[4],
n=num_repeats[4],
block=block,
)
)
channel_merge_layer = SimSPPF if not cspsppf else SimCSPSPPF
self.ERBlock_6 = nn.Sequential(
block(
in_channels=channels_list[4],
out_channels=channels_list[5],
kernel_size=3,
stride=2,
),
RepBlock(
in_channels=channels_list[5],
out_channels=channels_list[5],
n=num_repeats[5],
block=block,
),
channel_merge_layer(
in_channels=channels_list[5],
out_channels=channels_list[5],
kernel_size=5
)
)
def forward(self, x):
outputs = []
x = self.stem(x)
x = self.ERBlock_2(x)
if self.fuse_P2:
outputs.append(x)
x = self.ERBlock_3(x)
outputs.append(x)
x = self.ERBlock_4(x)
outputs.append(x)
x = self.ERBlock_5(x)
outputs.append(x)
x = self.ERBlock_6(x)
outputs.append(x)
return tuple(outputs)4.class CSPBepBackbone(nn.Module):
class CSPBepBackbone(nn.Module):
# CSPBepBackbone 模型。
"""
CSPBepBackbone module.
"""
def __init__(
self,
in_channels=3,
channels_list=None,
num_repeats=None,
block=RepVGGBlock,
csp_e=float(1)/2,
fuse_P2=False,
cspsppf=False,
stage_block_type="BepC3"
):
super().__init__()
assert channels_list is not None
assert num_repeats is not None
if stage_block_type == "BepC3":
stage_block = BepC3
elif stage_block_type == "MBLABlock":
stage_block = MBLABlock
else:
raise NotImplementedError
self.fuse_P2 = fuse_P2
self.stem = block(
in_channels=in_channels,
out_channels=channels_list[0],
kernel_size=3,
stride=2
)
self.ERBlock_2 = nn.Sequential(
block(
in_channels=channels_list[0],
out_channels=channels_list[1],
kernel_size=3,
stride=2
),
stage_block(
in_channels=channels_list[1],
out_channels=channels_list[1],
n=num_repeats[1],
e=csp_e,
block=block,
)
)
self.ERBlock_3 = nn.Sequential(
block(
in_channels=channels_list[1],
out_channels=channels_list[2],
kernel_size=3,
stride=2
),
stage_block(
in_channels=channels_list[2],
out_channels=channels_list[2],
n=num_repeats[2],
e=csp_e,
block=block,
)
)
self.ERBlock_4 = nn.Sequential(
block(
in_channels=channels_list[2],
out_channels=channels_list[3],
kernel_size=3,
stride=2
),
stage_block(
in_channels=channels_list[3],
out_channels=channels_list[3],
n=num_repeats[3],
e=csp_e,
block=block,
)
)
channel_merge_layer = SPPF if block == ConvBNSiLU else SimSPPF
if cspsppf:
channel_merge_layer = CSPSPPF if block == ConvBNSiLU else SimCSPSPPF
self.ERBlock_5 = nn.Sequential(
block(
in_channels=channels_list[3],
out_channels=channels_list[4],
kernel_size=3,
stride=2,
),
stage_block(
in_channels=channels_list[4],
out_channels=channels_list[4],
n=num_repeats[4],
e=csp_e,
block=block,
),
channel_merge_layer(
in_channels=channels_list[4],
out_channels=channels_list[4],
kernel_size=5
)
)
def forward(self, x):
outputs = []
x = self.stem(x)
x = self.ERBlock_2(x)
if self.fuse_P2:
outputs.append(x)
x = self.ERBlock_3(x)
outputs.append(x)
x = self.ERBlock_4(x)
outputs.append(x)
x = self.ERBlock_5(x)
outputs.append(x)
return tuple(outputs)5.class CSPBepBackbone_P6(nn.Module):
class CSPBepBackbone_P6(nn.Module):
# CSPBepBackbone+P6 模型。
"""
CSPBepBackbone+P6 module.
"""
def __init__(
self,
in_channels=3,
channels_list=None,
num_repeats=None,
block=RepVGGBlock,
csp_e=float(1)/2,
fuse_P2=False,
cspsppf=False
):
super().__init__()
assert channels_list is not None
assert num_repeats is not None
self.fuse_P2 = fuse_P2
self.stem = block(
in_channels=in_channels,
out_channels=channels_list[0],
kernel_size=3,
stride=2
)
self.ERBlock_2 = nn.Sequential(
block(
in_channels=channels_list[0],
out_channels=channels_list[1],
kernel_size=3,
stride=2
),
BepC3(
in_channels=channels_list[1],
out_channels=channels_list[1],
n=num_repeats[1],
e=csp_e,
block=block,
)
)
self.ERBlock_3 = nn.Sequential(
block(
in_channels=channels_list[1],
out_channels=channels_list[2],
kernel_size=3,
stride=2
),
BepC3(
in_channels=channels_list[2],
out_channels=channels_list[2],
n=num_repeats[2],
e=csp_e,
block=block,
)
)
self.ERBlock_4 = nn.Sequential(
block(
in_channels=channels_list[2],
out_channels=channels_list[3],
kernel_size=3,
stride=2
),
BepC3(
in_channels=channels_list[3],
out_channels=channels_list[3],
n=num_repeats[3],
e=csp_e,
block=block,
)
)
channel_merge_layer = SPPF if block == ConvBNSiLU else SimSPPF
if cspsppf:
channel_merge_layer = CSPSPPF if block == ConvBNSiLU else SimCSPSPPF
self.ERBlock_5 = nn.Sequential(
block(
in_channels=channels_list[3],
out_channels=channels_list[4],
kernel_size=3,
stride=2,
),
BepC3(
in_channels=channels_list[4],
out_channels=channels_list[4],
n=num_repeats[4],
e=csp_e,
block=block,
),
)
self.ERBlock_6 = nn.Sequential(
block(
in_channels=channels_list[4],
out_channels=channels_list[5],
kernel_size=3,
stride=2,
),
BepC3(
in_channels=channels_list[5],
out_channels=channels_list[5],
n=num_repeats[5],
e=csp_e,
block=block,
),
channel_merge_layer(
in_channels=channels_list[5],
out_channels=channels_list[5],
kernel_size=5
)
)
def forward(self, x):
outputs = []
x = self.stem(x)
x = self.ERBlock_2(x)
outputs.append(x)
x = self.ERBlock_3(x)
outputs.append(x)
x = self.ERBlock_4(x)
outputs.append(x)
x = self.ERBlock_5(x)
outputs.append(x)
x = self.ERBlock_6(x)
outputs.append(x)
return tuple(outputs)6.class Lite_EffiBackbone(nn.Module):
class Lite_EffiBackbone(nn.Module):
def __init__(self,
in_channels,
mid_channels,
out_channels,
num_repeat=[1, 3, 7, 3]
):
super().__init__()
out_channels[0]=24
self.conv_0 = ConvBNHS(in_channels=in_channels,
out_channels=out_channels[0],
kernel_size=3,
stride=2,
padding=1)
self.lite_effiblock_1 = self.build_block(num_repeat[0],
out_channels[0],
mid_channels[1],
out_channels[1])
self.lite_effiblock_2 = self.build_block(num_repeat[1],
out_channels[1],
mid_channels[2],
out_channels[2])
self.lite_effiblock_3 = self.build_block(num_repeat[2],
out_channels[2],
mid_channels[3],
out_channels[3])
self.lite_effiblock_4 = self.build_block(num_repeat[3],
out_channels[3],
mid_channels[4],
out_channels[4])
def forward(self, x):
outputs = []
x = self.conv_0(x)
x = self.lite_effiblock_1(x)
x = self.lite_effiblock_2(x)
outputs.append(x)
x = self.lite_effiblock_3(x)
outputs.append(x)
x = self.lite_effiblock_4(x)
outputs.append(x)
return tuple(outputs)
# @staticmethod
# @staticmethod 是一个python装饰器，用于标记一个静态方法。
# 静态方法是属于类而不是类的实例的方法。
# 静态方法不需要访问类的实例或实例的状态。因此，它们不接受 self 参数。
# 可以通过类名直接调用静态方法，也可以通过实例调用。
@staticmethod
# 它用于构建一个由多个网络模块组成的序列容器 nn.Sequential 。这个方法通常用于创建重复的网络结构，例如在构建深度学习模型时常用的残差块或卷积块。
# 1.num_repeat ：重复次数，指定了要重复多少次给定的网络结构。
# 2.in_channels ：输入通道数。
# 3.mid_channels ：中间通道数，通常用于中间层，以实现特征的转换和通道数的变化。
# 4.out_channels ：输出通道数。
def build_block(num_repeat, in_channels, mid_channels, out_channels):
# torch.nn.Sequential(*args)
# nn.Sequential 是 PyTorch 中的一个容器类，它用于按顺序包装一系列模块（如层、激活函数等），以便于构建顺序模型。在 nn.Sequential 容器中的每个模块都会按照它们被添加的顺序执行。
# 参数 ：
# *args ：可变数量的模块参数，可以是 nn.Module 的子类实例。
# modules ：一个有序字典或模块列表，包含要按顺序添加到 nn.Sequential 容器中的模块。
# 返回值 ：
# 返回一个 nn.Sequential 实例，它是一个 nn.Module 的子类，可以像普通模块一样被使用。
# 方法逻辑 ：
# 当你将模块添加到 nn.Sequential 容器中时，这些模块将按照它们被添加的顺序被包装起来。
# 在模型的前向传播（ forward 方法）中，输入数据将依次通过这些模块，每个模块的输出将成为下一个模块的输入。
# 初始化序列容器，创建一个空的 nn.Sequential 容器，用于按顺序存储网络模块。
block_list = nn.Sequential()
for i in range(num_repeat):
if i == 0:
block = Lite_EffiBlockS2(
in_channels=in_channels,
mid_channels=mid_channels,
out_channels=out_channels,
stride=2)
else:
block = Lite_EffiBlockS1(
in_channels=out_channels,
mid_channels=mid_channels,
out_channels=out_channels,
stride=1)
block_list.add_module(str(i), block)
# 返回一个包含重复网络模块的 nn.Sequential 容器。
return block_list