models\common.py
目录
common.py
1.所需的库和模块
2.def autopad(k, p=None):
3.class MP(nn.Module):
4.class SP(nn.Module):
5.class ReOrg(nn.Module):
6. class Concat(nn.Module):
7.class Chuncat(nn.Module):
8.class Shortcut(nn.Module):
9.class Foldcut(nn.Module):
10.class Conv(nn.Module):
11.class RobustConv(nn.Module):
12.class RobustConv2(nn.Module):
13.def DWConv(c1, c2, k=1, s=1, act=True):
14.class GhostConv(nn.Module):
15.class Stem(nn.Module):
16.class DownC(nn.Module):
17.class SPP(nn.Module):
18.class Bottleneck(nn.Module):
19.class Res(nn.Module):
20.class ResX(Res):
21.class Ghost(nn.Module):
22.class SPPCSPC(nn.Module):
23.class GhostSPPCSPC(SPPCSPC):
24.class GhostStem(Stem):
25.class BottleneckCSPA(nn.Module):
26.class BottleneckCSPB(nn.Module):
27.class BottleneckCSPC(nn.Module):
28.class ResCSPA(BottleneckCSPA):
29.class ResCSPB(BottleneckCSPB):
30.class ResCSPC(BottleneckCSPC):
31.class ResXCSPA(ResCSPA):
32.class ResXCSPB(ResCSPB):
33.class ResXCSPC(ResCSPC):
34.class GhostCSPA(BottleneckCSPA):
35.class GhostCSPB(BottleneckCSPB):
36.class GhostCSPC(BottleneckCSPC):
37.class ImplicitA(nn.Module):
38.class ImplicitM(nn.Module):
39.class RepConv(nn.Module):
40.class RepBottleneck(Bottleneck):
41.class RepBottleneckCSPA(BottleneckCSPA):
42.class RepBottleneckCSPB(BottleneckCSPB):
43.class RepBottleneckCSPC(BottleneckCSPC):
44.class RepRes(Res):
45.class RepResCSPA(ResCSPA):
46.class RepResCSPB(ResCSPB):
47.class RepResCSPC(ResCSPC):
48.class RepResX(ResX):
49.class RepResXCSPA(ResXCSPA):
50.class RepResXCSPB(ResXCSPB):
51.class RepResXCSPC(ResXCSPC):
52.class TransformerLayer(nn.Module):
53.class TransformerBlock(nn.Module):
54.class Focus(nn.Module):
55.class SPPF(nn.Module):
56.class Contract(nn.Module):
57.class Expand(nn.Module):
58.class NMS(nn.Module):
59.class autoShape(nn.Module):
60.class Detections:
61.class Classify(nn.Module):
62.def transI_fusebn(kernel, bn):
63.class ConvBN(nn.Module):
64.class OREPA_3x3_RepConv(nn.Module):
65.class RepConv_OREPA(nn.Module):
66.class WindowAttention(nn.Module):
67.class Mlp(nn.Module):
68.def window_partition(x, window_size):
69.def window_reverse(windows, window_size, H, W):
70.class SwinTransformerLayer(nn.Module):
71.class SwinTransformerBlock(nn.Module):
72.class STCSPA(nn.Module):
73.class STCSPB(nn.Module):
74.class STCSPC(nn.Module):
75.class WindowAttention_v2(nn.Module):
76.class Mlp_v2(nn.Module):
77.def window_partition_v2(x, window_size):
78.def window_reverse_v2(windows, window_size, H, W):
79.class SwinTransformerLayer_v2(nn.Module):
80.class SwinTransformer2Block(nn.Module):
81.class ST2CSPA(nn.Module):
82.class ST2CSPB(nn.Module):
83.class ST2CSPC(nn.Module):
import math
from copy import copy
from pathlib import Path
import numpy as np
import pandas as pd
import requests
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.ops import DeformConv2d
from PIL import Image
from torch.cuda import amp
from utils.datasets import letterbox
from utils.general import non_max_suppression, make_divisible, scale_coords, increment_path, xyxy2xywh
from utils.plots import color_list, plot_one_box
from utils.torch_utils import time_synchronized
##### basic ####
# autopad 函数用于自动计算卷积层的填充(padding),以实现“same”填充的效果,即输出特征图的大小与输入特征图的大小相同。这个函数的目的是确保卷积操作不会改变特征图的空间维度。
# 1.k :卷积核的大小(kernel size),可以是整数或列表/元组。
# 2.p :填充值(padding),默认为 None 。
def autopad(k, p=None): # kernel, padding
# 填充至“相同”。
# Pad to 'same'
# 如果 p 为 None ,则自动计算填充值。
if p is None:
# 如果 k 是整数,那么填充值 p 就是 k 除以 2 的结果,即 k // 2 。
# 如果 k 是列表或元组,那么对于 k 中的每个元素,计算其除以 2 的结果,并将这些结果作为列表返回。
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
# 返回计算得到的填充值 p 。
return p
# MP 类是一个简单的模块(Module),代表最大池化层(Max Pooling Layer)。这个类继承自 PyTorch 的 nn.Module ,用于在卷积神经网络中实现最大池化操作。
class MP(nn.Module):
# 1.k :池化核的大小,默认为 2。
def __init__(self, k=2):
# 调用父类 nn.Module 的构造函数,初始化模块。
super(MP, self).__init__()
# torch.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
# 对于输入信号的输入通道,提供2维最大池化(max pooling)操作。
# 如果输入的大小是(N,C,H,W),那么输出的大小是(N,C,H_out,W_out)。
# 参数:
# kernel_size(int or tuple) :max pooling的窗口大小。
# stride(int or tuple, optional) :max pooling的窗口移动的步长。默认值是kernel_size。
# padding(int or tuple, optional) :输入的每一条边补充0的层数。
# dilation(int or tuple, optional) :一个控制窗口中元素步幅的参数。
# return_indices :如果等于True,会返回输出最大值的序号,对于上采样操作会有帮助。
# ceil_mode :如果等于True,计算输出信号大小的时候,会使用向上取整,代替默认的向下取整的操作。
# 实例化一个 nn.MaxPool2d 对象,用于最大池化操作。 kernel_size=k 和 stride=k 指定了池化核的大小和步长。
self.m = nn.MaxPool2d(kernel_size=k, stride=k)
# x :输入的特征图。
def forward(self, x):
# self.m(x) :将最大池化操作应用于输入特征图 x 。
# 返回经过最大池化后的特征图。
return self.m(x)
# 输入特征图 x 的尺寸为 (1, 64, 256, 256) ,经过 MP 模块的最大池化操作后,输出特征图的尺寸变为 (1, 64, 128, 128) ,因为池化核大小和步长都是 2。
# SP 类代表的是一个步进最大池化层(Strided Max Pooling Layer),这是一个在卷积神经网络中用于下采样的层。这个类继承自 PyTorch 的 nn.Module ,用于实现带有步进的最大池化操作。
class SP(nn.Module):
# 1.k :池化核的大小,默认为 3。
# 2.s :步长,默认为 1。
def __init__(self, k=3, s=1):
# 调用父类 nn.Module 的构造函数,初始化模块。
super(SP, self).__init__()
# 实例化一个 nn.MaxPool2d 对象,用于最大池化操作。 kernel_size=k 指定了池化核的大小, stride=s 指定了步长, padding=k // 2 指定了填充,这里使用了自动填充( autopad )来确保输出特征图的大小与输入特征图相同(在步进为 1 时)。
self.m = nn.MaxPool2d(kernel_size=k, stride=s, padding=k // 2)
# 1.x :输入的特征图。
def forward(self, x):
# self.m(x) :将最大池化操作应用于输入特征图 x 。
# 返回经过最大池化后的特征图。
return self.m(x)
# 输入特征图 x 的尺寸为 (1, 64, 256, 256) ,经过 SP 模块的最大池化操作后,输出特征图的尺寸保持不变为 (1, 64, 256, 256) ,因为步长 s 设置为 1,且填充 padding 被设置为 k // 2 ,以保持特征图的空间维度。
# ReOrg 类是一个 PyTorch 神经网络模块,它的目的是对输入的特征图进行重组(ReOrganization),以实现一种特定的空间降采样和通道升维的效果。这种操作在某些神经网络架构中被用来减少特征图的空间尺寸同时增加通道数,有助于特征的融合和信息的传递。
class ReOrg(nn.Module):
# 初始化 ReOrg 类,不包含任何额外的参数。
def __init__(self):
super(ReOrg, self).__init__()
# 1.x :输入的特征图,其形状为 (b, c, w, h) ,其中 b 是批次大小, c 是通道数, w 和 h 分别是特征图的宽度和高度。
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
# x[..., ::2, ::2] :从输入特征图中选取每两个像素中的第一个像素,即每隔一个像素取一个,这会将特征图的高度和宽度都减少到原来的一半。
# x[..., 1::2, ::2] :从输入特征图中选取每两个像素中的第二个像素(从1开始计数),同样减少特征图的高度和宽度到一半。
# x[..., ::2, 1::2] :从输入特征图中选取每隔一行的第一个像素,每隔一列的第二个像素,减少特征图的高度和宽度到一半。
# x[..., 1::2, 1::2] :从输入特征图中选取每隔一行的第二个像素,每隔一列的第二个像素,减少特征图的高度和宽度到一半。
# torch.cat([...], 1) :将上述四个结果沿着通道维度(即维度1)进行拼接,得到最终的输出特征图,其形状为 (b, 4c, w/2, h/2) 。
return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
# 输入特征图 x 的尺寸为 (1, 8, 256, 256) ,经过 ReOrg 模块的特征重组操作后,输出特征图的尺寸变为 (1, 32, 128, 128) ,其中通道数增加了4倍,空间尺寸减少了一半。这种操作有助于在保持特征图信息的同时减少计算量。
# Concat 类是一个 PyTorch 神经网络模块,用于将多个张量(tensors)沿着指定的维度进行拼接。这个类继承自 nn.Module ,提供了一个灵活的方式来合并特征图或其他张量,这在构建复杂的神经网络结构时非常有用。
class Concat(nn.Module):
# 1.dimension :指定拼接张量的维度,默认为 1,即通道维度(对于图像数据,通常是通道维度)。
def __init__(self, dimension=1):
# 调用父类 nn.Module 的构造函数,初始化模块。
super(Concat, self).__init__()
# 存储传入的维度参数,用于后续的张量拼接。
self.d = dimension
# 1.x :一个包含多个张量的元组或列表,这些张量将被拼接。
def forward(self, x):
# 使用 PyTorch 的 torch.cat 函数将传入的张量列表 x 沿着 self.d 指定的维度进行拼接。
# 返回拼接后的张量。
return torch.cat(x, self.d)
# Chuncat 类是一个 PyTorch 神经网络模块,它首先将输入的每个张量沿着指定维度分割成两部分,然后将这两部分分别拼接在一起。这个操作可以看作是一种特殊的特征重组,其中每个输入张量被分成两半,然后这两半被交叉拼接。
class Chuncat(nn.Module):
# 1.dimension :指定分割和拼接张量的维度,默认为 1,即通道维度。
def __init__(self, dimension=1):
# 调用父类 nn.Module 的构造函数,初始化模块。
super(Chuncat, self).__init__()
# 存储传入的维度参数,用于后续的张量分割和拼接。
self.d = dimension
# 1.x :一个包含多个张量的元组或列表,这些张量将被分割和拼接。
def forward(self, x):
# 初始化两个空列表,用于存储分割后的张量的两半。
x1 = []
x2 = []
# 遍历输入的张量列表 x ,其中每个 xi 是一个张量。
for xi in x:
# torch.chunk(input, chunks, dim=0)
# torch.chunk 是 PyTorch 中的一个函数,它将张量(tensor)分割成指定数量的块(chunks)。每个块在指定的维度上具有相等的大小。如果张量不能被均匀分割,则最后一个块可能会比其他块小。
# 参数 :
# input :要被分割的输入张量。
# chunks :一个整数,表示要将输入张量分割成多少块。
# dim :一个整数,指定沿着哪个维度进行分割。默认是0,即第一个维度。
# 返回值 :
# 返回一个包含分割后块的元组,每个块都是一个张量。
# 使用 chunk 方法将张量 xi 沿着维度 self.d 分成两个相等的部分。 chunk 方法将张量分割成指定数量的块,这里指定的是2,所以张量被分成两半。
xi1, xi2 = xi.chunk(2, self.d)
# 将分割得到的两半分别添加到列表 x1 和 x2 中。
x1.append(xi1)
x2.append(xi2)
# 使用 torch.cat 函数将列表 x1 和 x2 中的张量沿着维度 self.d 合并。 x1+x2 将两个列表中的元素交错排列,实现交错合并的效果。
return torch.cat(x1+x2, self.d)
# Shortcut 类是一个 PyTorch 神经网络模块,用于实现快捷连接(shortcut connection)或残差连接(residual connection)。这种连接通常用于残差网络(ResNet)和其他需要将输入直接添加到后续层的输出上的架构中。
class Shortcut(nn.Module):
# 1.dimension :指定连接的维度,默认为 0,即第一个维度(批次维度)。
def __init__(self, dimension=0):
super(Shortcut, self).__init__()
# 存储传入的维度参数,虽然在这个简单的 Shortcut 类中未使用,但可以用于更复杂的快捷连接实现,例如在不同维度上进行连接。
self.d = dimension
# 1.x :一个包含两个张量的元组或列表,这两个张量将被相加。
def forward(self, x):
# x[0] + x[1] :将两个张量相加并返回结果。
# 返回相加后的张量。
return x[0]+x[1]
# 示例中, feature_map1 和 feature_map2 被相加,输出特征图的通道数保持不变,空间尺寸也保持不变。
# Shortcut 类提供了一个简单的方式来实现快捷连接,这在许多现代神经网络架构中是一个常见的模式,有助于缓解梯度消失问题并提高训练深度网络的效率。
# Foldcut 类是一个 PyTorch 神经网络模块,它实现了一个类似于 Shortcut 的操作,但具有更复杂的处理方式。
# 这个类将输入张量沿着指定维度分割成两部分,然后将这两部分相加。这种操作可以看作是一种特殊的特征重组,其中输入张量被分成两半,然后这两半被相加。
class Foldcut(nn.Module):
# 1.dimension :指定分割和相加张量的维度,默认为 0,即第一个维度(批次维度)。
def __init__(self, dimension=0):
super(Foldcut, self).__init__()
self.d = dimension
# 1.x :输入的张量。
def forward(self, x):
# 使用 PyTorch 的 chunk 方法将输入张量 x 沿着 self.d 指定的维度分割成两个相等的部分。这里指定的是2,所以张量被分成两半。
x1, x2 = x.chunk(2, self.d)
# 将分割得到的两半相加。
# 返回相加后的张量。
return x1+x2
# 这个 Conv 类是 PyTorch 中的一个自定义模块,用于实现一个标准的卷积层,包括卷积操作、批量归一化(Batch Normalization)和激活函数(通常是 LeakyReLU)。这个类提供了灵活性,允许用户指定卷积核大小、步长、填充、组数等参数,并可以选择是否使用激活函数。
class Conv(nn.Module):
# 标准卷积。
# Standard convolution
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.k :卷积核大小,默认为 1。
# 4.s :步长,默认为 1。
# 5.p :填充,如果为 None ,则自动计算填充以实现“same”卷积。
# 6.g :组数,默认为 1,用于分组卷积。
# 7.act :是否使用激活函数,默认为 True ,使用 LeakyReLU。如果 act 是 nn.Module 的实例,则使用该模块作为激活函数;如果 act 是 False ,则不使用激活函数,使用 nn.Identity() 。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super(Conv, self).__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
self.bn = nn.BatchNorm2d(c2)
# nn.Identity()
# 是 PyTorch 中的一个层(layer)。它实际上是一个恒等映射,不对输入进行任何变换或操作,只是简单地将输入返回作为输出。
self.act = nn.LeakyReLU(0.1, inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
# 标准的前向传播过程,包括卷积、批量归一化和激活函数。
def forward(self, x):
return self.act(self.bn(self.conv(x)))
# 融合前向传播,直接应用卷积和激活函数,跳过批量归一化。这在某些情况下用于模型剪枝或加速推理。
def fuseforward(self, x):
return self.act(self.conv(x))
# RobustConv 类是一个自定义的 PyTorch 模块,它实现了一种更健壮的卷积操作,通常用于需要高核大小(例如7到11)的场景,比如下采样或其他层。这种卷积通常在训练300到450个epoch后能够表现出更好的性能和鲁棒性。
class RobustConv(nn.Module):
# 健壮卷积(使用高内核大小 7-11 进行:下采样和其他层)。训练 300 - 450 个时期。
# Robust convolution (use high kernel size 7-11 for: downsampling and other layers). Train for 300 - 450 epochs.
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.k :卷积核大小,默认为7。
# 4.s :步长,默认为1。
# 5.p :填充,如果为 None ,则自动计算填充以实现“same”卷积。
# 6.g :组数,默认为1,用于分组卷积。
# 7.act :是否使用激活函数,默认为 True 。
# 8.layer_scale_init_value :层缩放初始化值,默认为1e-6。
def __init__(self, c1, c2, k=7, s=1, p=None, g=1, act=True, layer_scale_init_value=1e-6): # ch_in, ch_out, kernel, stride, padding, groups
super(RobustConv, self).__init__()
# 深度可分离卷积层,使用 Conv 类实现,其中 g=c1 表示每个输入通道都有自己的卷积核,实现深度卷积。
self.conv_dw = Conv(c1, c1, k=k, s=s, p=p, g=c1, act=act)
# 1x1卷积层,用于进一步减少或增加通道数。
self.conv1x1 = nn.Conv2d(c1, c2, 1, 1, 0, groups=1, bias=True)
# 一个可学习的参数,用于层缩放(Layer Scale),如果 layer_scale_init_value 大于0,则初始化为该值,否则为 None 。
self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(c2)) if layer_scale_init_value > 0 else None
def forward(self, x):
# 将输入张量 x 转换为通道最后(channels_last)的内存格式。这种格式对于某些硬件(如GPU)来说,可以提高内存访问效率和卷积操作的性能。
x = x.to(memory_format=torch.channels_last)
# 将输入张量 x 通过深度可分离卷积层 self.conv_dw 。这种卷积首先在深度方向上进行卷积(每个输入通道一个卷积核),然后在空间方向上进行卷积(所有深度方向的输出合并使用同一个卷积核)。
# 将深度可分离卷积的输出通过1x1卷积层 self.conv1x1 。1x1卷积可以改变通道数,同时引入不同通道之间的信息混合。
x = self.conv1x1(self.conv_dw(x))
# 检查是否存在可学习的参数 self.gamma 。
if self.gamma is not None:
# torch.Tensor.mul(input, other)
# 参数 :
# input :要进行乘法操作的张量。
# other :另一个张量,可以是标量或与 input 形状相同的张量。
# 描述 :
# 如果 other 是一个标量,则 input 中的每个元素都会乘以这个标量。 如果 other 是一个张量,那么 input 和 other 必须具有相同的形状,或者 other 可以是广播(broadcast)到 input 形状的张量。在这种情况下, input 和 other 会逐元素相乘。
# 返回值 :
# 返回修改后的 input 张量。
# 如果 self.gamma 存在,将其重塑为适合广播的格式,并与输出 x 相乘,实现层缩放。这种缩放可以增强或抑制特定通道的特征响应。
x = x.mul(self.gamma.reshape(1, -1, 1, 1))
return x
# RobustConv2 类是一个 PyTorch 神经网络模块,它结合了标准的卷积操作和转置卷积(也称为反卷积)操作,用于实现一种更复杂的特征融合和上采样。这种结构特别适用于需要同时考虑下采样和上采样的场景,例如在 CSP(Cross Stage Partial)网络中。
class RobustConv2(nn.Module):
# 健壮卷积 2(对 CSP 中的一条路径使用 [32, 5, 2] 或 [32, 7, 4] 或 [32, 11, 8])。
# Robust convolution 2 (use [32, 5, 2] or [32, 7, 4] or [32, 11, 8] for one of the paths in CSP).
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.k :卷积核大小,默认为7。
# 4.s :步长,默认为4。
# 5.p :填充,如果为 None ,则自动计算填充以实现“same”卷积。
# 6.g :组数,默认为1,用于分组卷积。
# 7.act :是否使用激活函数,默认为 True 。
# 8.layer_scale_init_value :层缩放初始化值,默认为1e-6。
def __init__(self, c1, c2, k=7, s=4, p=None, g=1, act=True, layer_scale_init_value=1e-6): # ch_in, ch_out, kernel, stride, padding, groups
super(RobustConv2, self).__init__()
# 深度可分离卷积层,使用 Conv 类实现,其中步长 s 用于下采样。
self.conv_strided = Conv(c1, c1, k=k, s=s, p=p, g=c1, act=act)
# 转置卷积层,用于上采样,其步长 s 与下采样步长相对应,以实现特征图的空间尺寸恢复。
self.conv_deconv = nn.ConvTranspose2d(in_channels=c1, out_channels=c2, kernel_size=s, stride=s,
padding=0, bias=True, dilation=1, groups=1
)
# 一个可学习的参数,用于层缩放(Layer Scale),如果 layer_scale_init_value 大于0,则初始化为该值,否则为 None 。
self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(c2)) if layer_scale_init_value > 0 else None
def forward(self, x):
# 首先通过 self.conv_strided 进行下采样,然后通过 self.conv_deconv 进行上采样。
x = self.conv_deconv(self.conv_strided(x))
# 检查是否存在可学习的参数 self.gamma 。
if self.gamma is not None:
# 如果 self.gamma 存在,将其重塑为适合广播的格式,并与输出 x 相乘,实现层缩放。
x = x.mul(self.gamma.reshape(1, -1, 1, 1))
return x
# 输入特征图 x 的尺寸为 (1, 16, 256, 256) ,经过 RobustConv2 模块的卷积操作后,输出特征图的尺寸变为 (1, 32, 64, 64) ,通道数增加到32,空间尺寸减半然后通过转置卷积恢复。这种结构可以在保持特征图信息的同时,实现有效的特征融合和尺寸调整。
# DWConv 函数是一个深度可分离卷积(Depthwise Separable Convolution)的封装。这种卷积是一种优化技术,用于减少模型的计算量和参数数量,同时保持相似的性能。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数。
# 3.k :卷积核的大小,默认为1。
# 4.s :步长,默认为1。
# 5.act :是否应用激活函数,默认为True。
def DWConv(c1, c2, k=1, s=1, act=True):
# 深度卷积。
# Depthwise convolution
# 函数内部调用了 Conv 函数,这是一个通用的卷积层函数。 DWConv 通过设置 g 参数为输入通道数和输出通道数的最大公约数(GCD),来实现深度可分离卷积的效果。
# g 参数在这里代表分组卷积的组数,当 g 等于输入通道数时,每个输入通道都有一个独立的卷积核,这就是深度卷积(Depthwise Convolution)。
# 随后, Conv 函数会应用逐点卷积(Pointwise Convolution),即将深度卷积的输出通道合并为所需的输出通道数。
# 深度可分离卷积由两部分组成 :
# 1. 深度卷积(Depthwise Convolution) :每个输入通道独立地应用一个卷积核,输出与输入通道数相同数量的特征图。
# 2. 逐点卷积(Pointwise Convolution) :将深度卷积的输出通道通过1x1卷积核合并为所需的输出通道数。
return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
这段代码定义了一个名为 GhostConv 的类,它是PyTorch中的一个模块( nn.Module ),用于实现GhostNet中的Ghost Convolution。Ghost Convolution是一种高效的卷积操作,旨在减少模型的参数数量和计算量,同时保持性能。
class GhostConv(nn.Module):
# Ghost Convolution https://github.com/huawei-noah/ghostnet
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 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): # ch_in, ch_out, kernel, stride, groups
super(GhostConv, self).__init__()
c_ = c2 // 2 # hidden channels
# 这是一个普通的卷积层,用于将输入通道 c1 转换为中间通道 c_ 。中间通道数 c_ 是输出通道数 c2 的一半。
self.cv1 = Conv(c1, c_, k, s, None, g, act)
# 这是第二个卷积层,它将 self.cv1 的输出(即中间通道 c_ )进一步处理,使用5x5的卷积核和步长1。
self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
def forward(self, x):
# 输入数据 x 首先通过 self.cv1 进行卷积操作,得到中间特征 y 。
y = self.cv1(x)
# 然后, y 被传递给 self.cv2 进行进一步的卷积操作。
# 最后, self.cv1 和 self.cv2 的输出通过 torch.cat 函数在通道维度(维度1)上进行拼接,得到最终的输出。
return torch.cat([y, self.cv2(y)], 1)
# Ghost Convolution 的特点 :
# Ghost Convolution 的核心思想是通过在卷积操作中引入廉价的辅助通道(cheap auxiliary channels),这些通道通过在原始特征图上应用不同的随机或可学习的稀疏掩码来生成。
# 这种方法可以在不显著增加参数数量和计算量的情况下,增加网络的感受野和表达能力。
# 在代码中, self.cv1 和 self.cv2 分别对应于Ghost Convolution中的两个主要步骤 :
# 第一步是普通的卷积操作,第二步是对第一步的输出应用更大的卷积核,以增加感受野。最后,这两个步骤的输出被拼接在一起,以实现Ghost Convolution的效果。这种方法在保持模型性能的同时,显著减少了模型的参数数量和计算量。
# 这段代码定义了一个名为 Stem 的类,它是PyTorch中的一个模块( nn.Module ),用于构建神经网络的Stem部分,通常用于网络的初始层,以高效地从输入数据中提取特征。Stem层通常用于代替传统的卷积层,以减少参数数量和计算量,同时保持或提升特征提取的能力。
class Stem(nn.Module):
# Stem 茎。
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 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): # ch_in, ch_out, kernel, stride, padding, groups
super(Stem, self).__init__()
c_ = int(c2/2) # hidden channels
# 在初始化方法中, Stem 类创建了四个卷积层和一个最大池化层。
# 第一个卷积层,用于将输入通道 c1 转换为中间通道 c_ ,使用3x3的卷积核和步长2。
self.cv1 = Conv(c1, c_, 3, 2)
# 第二个卷积层,用于进一步处理 self.cv1 的输出,使用1x1的卷积核和步长1。
self.cv2 = Conv(c_, c_, 1, 1)
# 第三个卷积层,用于进一步处理 self.cv2 的输出,使用3x3的卷积核和步长2。
self.cv3 = Conv(c_, c_, 3, 2)
# torch.nn.MaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False)
# 对于输入信号的输入通道,提供2维最大池化(max pooling)操作。
# 如果输入的大小是(N,C,H,W),那么输出的大小是(N,C,H_out,W_out)。
# 参数:
# kernel_size(int or tuple) :max pooling的窗口大小。
# stride(int or tuple, optional) :max pooling的窗口移动的步长。默认值是kernel_size。
# padding(int or tuple, optional) :输入的每一条边补充0的层数。
# dilation(int or tuple, optional) :一个控制窗口中元素步幅的参数。
# return_indices :如果等于True,会返回输出最大值的序号,对于上采样操作会有帮助。
# ceil_mode :如果等于True,计算输出信号大小的时候,会使用向上取整,代替默认的向下取整的操作。
# 最大池化层,用于下采样,使用2x2的池化核和步长2。
self.pool = torch.nn.MaxPool2d(2, stride=2)
# 第四个卷积层,用于将 self.cv3 和 self.pool 的输出合并,并将通道数从 2 * c_ 转换为最终的输出通道数 c2 。
self.cv4 = Conv(2 * c_, c2, 1, 1)
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作。
x = self.cv1(x)
# 然后, self.cv2 的输出被传递给 self.cv3 进行进一步的卷积操作。同时,原始输入 x 也被传递给 self.pool 进行最大池化操作。
# 最后, self.cv3 的输出和 self.pool 的输出在通道维度(维度1)上进行拼接,并通过 self.cv4 进行最终的卷积操作,得到最终的输出。
return self.cv4(torch.cat((self.cv3(self.cv2(x)), self.pool(x)), dim=1))
# Stem 层的特点 :
# Stem 层的设计通常是为了在网络的开始阶段就进行有效的特征提取和下采样。通过使用不同大小的卷积核和池化操作, Stem 层可以捕获不同尺度的特征,并减少后续层的计算量。
# 在代码中, Stem 类通过一系列的卷积和池化操作,实现了这一目的,同时通过拼接不同来源的特征,增强了网络的特征融合能力。
# 这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型的初始化阶段。
# 这段代码定义了一个名为 DownC 的类,它是PyTorch中的一个模块( nn.Module ),用于实现一个下采样模块,这个模块在结构上类似于YOLOv3-SPP(Spatial Pyramid Pooling)中的空间金字塔池化层。这种结构通常用于在神经网络中进行特征的下采样和特征融合。
class DownC(nn.Module):
# YOLOv3-SPP 中使用的空间金字塔池化层。
# Spatial pyramid pooling layer used in YOLOv3-SPP
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out ),在这里输出通道数实际上是 c2 的一半,因为最后两个分支的输出会在通道维度上进行拼接。
# 3.n :这个参数在代码中没有使用,可能是用于控制某些操作的数量,但在当前实现中被忽略了。
# 4.k :步长,用于卷积和池化操作。
def __init__(self, c1, c2, n=1, k=2):
super(DownC, self).__init__()
c_ = int(c1) # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将 self.cv1 的输出进一步转换为 c2//2 通道,使用3x3的卷积核和步长 k 。
self.cv2 = Conv(c_, c2//2, 3, k)
# 第三个卷积层,用于将输入通道 c1 直接转换为 c2//2 通道,使用1x1的卷积核和步长1。
self.cv3 = Conv(c1, c2//2, 1, 1)
# 最大池化层,用于下采样,使用 k 大小的池化核和步长 k 。
self.mp = nn.MaxPool2d(kernel_size=k, stride=k)
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作,然后 self.cv2 的输出是第一个分支的特征。
# 同时, x 也被传递给 self.mp 进行最大池化操作,然后 self.cv3 的输出是第二个分支的特征。
# 最后,这两个分支的输出在通道维度(维度1)上进行拼接,得到最终的输出。
return torch.cat((self.cv2(self.cv1(x)), self.cv3(self.mp(x))), dim=1)
# 下采样模块的特点 :
# 这个 DownC 模块的设计意图是在保持特征空间分辨率的同时,增加特征的通道数。通过使用不同大小的卷积核和池化操作,这个模块可以捕获不同尺度的特征,并在通道维度上进行融合。这种方法有助于在下采样的同时保留更多的空间信息,并且通过特征融合增强了网络的特征表达能力。
# 在代码中, DownC 类通过一系列的卷积和池化操作,实现了这一目的,同时通过拼接不同来源的特征,增强了网络的特征融合能力。这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型中的下采样操作。
# 这段代码定义了一个名为 SPP 的类,它是PyTorch中的一个模块( nn.Module ),用于实现YOLOv3-SPP中的空间金字塔池化(Spatial Pyramid Pooling, SPP)层。SPP层是一种有效的特征提取结构,它可以在不同尺度上捕获图像特征,增强模型对不同尺度物体的识别能力。
class SPP(nn.Module):
# YOLOv3-SPP 中使用的空间金字塔池化层。
# Spatial pyramid pooling layer used in YOLOv3-SPP
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.k :一个元组,包含用于不同池化层的核大小,默认为(5, 9, 13)。
def __init__(self, c1, c2, k=(5, 9, 13)):
super(SPP, self).__init__()
c_ = c1 // 2 # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将经过池化层处理后的合并特征转换为输出通道 c2 ,使用1x1的卷积核和步长1。
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
# nn.ModuleList()
# nn.ModuleList 是 PyTorch 中的一个容器类,它继承自 nn.Module ,用于存储一个有序的模块列表。 nn.ModuleList 的主要作用是保持子模块的有序性,并且确保在模型的参数更新时,这些子模块的参数也会被包括在内。
# class ModuleList(nn.Module): def __init__(self, modules=None):
# module_list = nn.ModuleList() module_list.append(nn.Linear(2, 3)) module_list.append(nn.ReLU())
# nn.ModuleList 的关键特性是它会自动注册其包含的模块为子模块,这意味着当你调用 model.parameters() 或 model.named_parameters() 时, ModuleList 中的模块参数也会被包括在内,从而在模型训练时会被优化。
# 此外, nn.ModuleList 支持索引、切片、长度查询等操作,使得管理模块集合变得非常方便。
# 一个模块列表,包含不同核大小的最大池化层,每个池化层的核大小由 k 参数指定。
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作。
x = self.cv1(x)
# 然后通过 self.m 中的每个最大池化层进行池化操作。每个池化层都会输出一个特征图,这些特征图与原始特征图 x 一起在通道维度(维度1)上进行拼接。最后,拼接后的特征图通过 self.cv2 进行卷积操作,得到最终的输出。
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
# 空间金字塔池化层的特点 :
# SP P层的主要特点是能够在不同尺度上捕获图像特征,这对于目标检测等任务非常重要,因为它可以帮助模型识别不同大小的物体。通过使用不同核大小的池化层,SPP层可以模拟不同大小的卷积核,从而在不增加额外参数的情况下捕获多尺度特征。
# 在代码中, SPP 类通过一系列的池化操作和卷积操作,实现了这一目的。这种方法在保持模型性能的同时,增强了模型对不同尺度特征的捕获能力,适合用于深度学习模型中的特征提取阶段。
# 这段代码定义了一个名为 Bottleneck 的类,它是PyTorch中的一个模块( nn.Module ),用于实现Darknet架构中的瓶颈(bottleneck)结构。
# 瓶颈结构是一种常见的网络设计模式,用于减少网络的参数数量和计算量,同时保持或提升网络的性能。这种结构通常包含两个主要的卷积层,中间可能有一个或多个卷积层。
class Bottleneck(nn.Module):
# Darknet bottleneck
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.shortcut :是否使用快捷连接(shortcut connection),默认为True。
# 4.g :分组卷积的组数,默认为1。
# 5.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
super(Bottleneck, self).__init__()
c_ = int(c2 * e) # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将 self.cv1 的输出进一步转换为输出通道 c2 ,使用3x3的卷积核和步长1,以及分组数 g 。
self.cv2 = Conv(c_, c2, 3, 1, g=g)
# self.add 是一个布尔值,用于指示是否应该将 self.cv2(self.cv1(x)) 的结果与输入 x 相加。这个快捷连接(shortcut connection)只有在 shortcut 为True且输入通道数 c1 等于输出通道数 c2 时才启用。
self.add = shortcut and c1 == c2
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作,然后 self.cv2 的输出根据 self.add 的值决定是否与输入 x 相加。
# 如果 self.add 为True,则执行残差连接,即将 self.cv2(self.cv1(x)) 的结果与输入 x 相加;否则,直接返回 self.cv2(self.cv1(x)) 的结果。
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
# 瓶颈结构的特点 :
# 瓶颈结构的主要优点是减少了网络的参数数量和计算量,同时通过残差连接保持了网络的性能。这种结构在深度学习中非常流行,特别是在卷积神经网络(CNN)中。通过使用1x1卷积减少通道数,然后通过3x3卷积提取特征,最后再通过1x1卷积增加通道数,瓶颈结构能够有效地捕获特征并减少计算量。
# 在代码中, Bottleneck 类通过这种结构实现了特征提取和通道数的调整,同时通过残差连接保持了网络的性能。这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型中的中间层。
# 这段代码定义了一个名为 Res 的类,它是PyTorch中的一个模块( nn.Module ),用于实现ResNet架构中的瓶颈(bottleneck)结构。ResNet(残差网络)是一种深度学习架构,它通过引入残差连接(shortcut connections)来解决深度网络中的梯度消失问题,使得可以训练更深的网络。
class Res(nn.Module):
# ResNet bottleneck
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.shortcut :是否使用快捷连接(shortcut connection),默认为True。
# 4.g :分组卷积的组数,默认为1。
# 5.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
super(Res, self).__init__()
c_ = int(c2 * e) # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将 self.cv1 的输出进一步处理,仍然保持通道数为 c_ ,使用3x3的卷积核和步长1,以及分组数 g 。
self.cv2 = Conv(c_, c_, 3, 1, g=g)
# 第三个卷积层,用于将 self.cv2 的输出转换为输出通道 c2 ,使用1x1的卷积核和步长1。
self.cv3 = Conv(c_, c2, 1, 1)
# self.add 是一个布尔值,用于指示是否应该将 self.cv3(self.cv2(self.cv1(x))) 的结果与输入 x 相加。这个快捷连接(shortcut connection)只有在 shortcut 为True且输入通道数 c1 等于输出通道数 c2 时才启用。
self.add = shortcut and c1 == c2
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作,然后 self.cv2 的输出再通过 self.cv3 进行卷积操作。
# 根据 self.add 的值决定是否将这个结果与输入 x 相加。
# 如果 self.add 为True,则执行残差连接,即将 self.cv3(self.cv2(self.cv1(x))) 的结果与输入 x 相加;否则,直接返回 self.cv3(self.cv2(self.cv1(x))) 的结果。
return x + self.cv3(self.cv2(self.cv1(x))) if self.add else self.cv3(self.cv2(self.cv1(x)))
# 残差连接的特点 :
# 残差连接是 ResNet 架构的核心特点,它允许网络中的每个层直接学习到输入的残差(即输入与输出之间的差异),而不是直接学习到输出。这样做的好处是可以缓解梯度消失问题,使得网络可以更深,同时保持训练的稳定性。
# 在代码中, Res 类通过这种结构实现了特征提取和通道数的调整,同时通过残差连接保持了网络的性能。这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型中的中间层。
# 这段代码定义了一个名为 ResX 的类,它是 Res 类的子类。 ResX 类继承了 Res 类的所有属性和方法,并在其基础上进行了一些修改。这个类看起来是为了实现一个特定的变种的ResNet瓶颈结构,其中可能包含了一些额外的特性或者参数的调整。
class ResX(Res):
# ResNet bottleneck
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.shortcut :是否使用快捷连接(shortcut connection),默认为True。
# 4.g :分组卷积的组数,默认为32。
# 5.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
def __init__(self, c1, c2, shortcut=True, g=32, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
# 在 ResX 类的初始化方法中,首先调用了父类 Res 的初始化方法,传入了相同的参数。
super().__init__(c1, c2, shortcut, g, e)
# 然后, ResX 类重新计算了隐藏通道数 c_ ,这是基于输出通道数 c2 和扩张率 e 的乘积的一半。
c_ = int(c2 * e) # hidden channels
# 特点 :
# 分组卷积 :ResX 类在初始化方法中将分组卷积的组数 g 设置为32,这是一种常见的做法,用于在卷积操作中提高效率和效果。分组卷积可以减少模型的参数数量和计算量,同时保持或提升模型的性能。
# 扩张率 :扩张率 e 用于控制隐藏通道数 c_ ,这个参数可以调整网络的容量,即网络能够学习的特征的复杂度。
# 继承和扩展 :
# ResX 类继承自 Res 类,这意味着它直接继承了 Res 类的所有属性和方法,包括三个卷积层( cv1 、 cv2 、 cv3 )和一个快捷连接的逻辑( self.add )。通过继承, ResX 可以重用 Res 类的代码,同时可以添加或修改特定的功能。
# 这段代码定义了一个名为 Ghost 的类,它是PyTorch中的一个模块( nn.Module ),用于实现GhostNet架构中的Ghost Bottleneck。GhostNet是一种轻量级的深度学习模型,由华为诺亚方舟实验室开发,旨在减少模型的参数数量和计算量,同时保持较高的性能。
class Ghost(nn.Module):
# Ghost Bottleneck https://github.com/huawei-noah/ghostnet
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.k :卷积核大小,默认为3。
# 4.s :步长,默认为1。
def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
super(Ghost, self).__init__()
c_ = c2 // 2
# self.conv :这是一个包含三个子模块的序列。
# 第一个子模块是 GhostConv ,用于执行1x1的卷积操作,将输入通道 c1 转换为中间通道 c_ 。
# 第二个子模块是 DWConv ,用于执行深度可分离卷积(Depthwise Convolution),如果步长 s 为2,则使用 DWConv ,否则使用 nn.Identity() ,后者相当于一个恒等映射。
# 第三个子模块是另一个 GhostConv ,用于执行1x1的卷积操作,将中间通道 c_ 转换为输出通道 c2 ,且不应用激活函数。
self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
# self.shortcut :这是一个快捷连接,用于将输入 x 直接映射到输出通道 c2 。如果步长 s 为2,则包含一个 DWConv 和一个 Conv ,用于调整通道数和应用恒等映射;否则,使用 nn.Identity() ,后者相当于一个恒等映射。
self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.conv 进行卷积操作,然后通过 self.shortcut 进行快捷连接操作。最后,将这两个结果相加,得到最终的输出。
return self.conv(x) + self.shortcut(x)
# Ghost Bottleneck 的特点 :
# Ghost Bottleneck 的核心思想是通过引入廉价的辅助通道(cheap auxiliary channels),这些通道通过在原始特征图上应用不同的随机或可学习的稀疏掩码来生成。这种方法可以在不显著增加参数数量和计算量的情况下,增加网络的感受野和表达能力。
# 在代码中, Ghost 类通过 GhostConv 和 DWConv 的组合实现了 Ghost Bottleneck 的效果,同时通过快捷连接保持了网络的性能。这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型中的中间层。
##### end of basic #####
##### cspnet #####
# 这段代码定义了一个名为 SPPCSPC 的类,它是PyTorch中的一个模块( nn.Module ),用于实现Cross Stage Partial Networks (CSP) 结合空间金字塔池化(SPP)的结构。CSP是一种网络结构,旨在提高网络的效率和性能,而SPP是一种用于捕获多尺度特征的技术。
class SPPCSPC(nn.Module):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.n :这个参数在代码中没有使用,可能是用于控制某些操作的数量,但在当前实现中被忽略了。
# 4.shortcut :是否使用快捷连接,默认为False。
# 5.g :分组卷积的组数,默认为1。
# 6.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
# 7.k :一个元组,包含用于不同池化层的核大小,默认为(5, 9, 13)。
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)):
super(SPPCSPC, self).__init__()
c_ = int(2 * c2 * e) # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv2 = Conv(c1, c_, 1, 1)
# 第三个卷积层,用于将 self.cv1 的输出进一步转换为隐藏通道 c_ ,使用3x3的卷积核和步长1。
self.cv3 = Conv(c_, c_, 3, 1)
# 第四个卷积层,用于将 self.cv3 的输出进一步转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv4 = Conv(c_, c_, 1, 1)
# 一个模块列表,包含不同核大小的最大池化层,每个池化层的核大小由 k 参数指定。
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
# 第五个卷积层,用于将经过池化层处理后的合并特征转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv5 = Conv(4 * c_, c_, 1, 1)
# 第六个卷积层,用于将 self.cv5 的输出进一步转换为隐藏通道 c_ ,使用3x3的卷积核和步长1。
self.cv6 = Conv(c_, c_, 3, 1)
# 第七个卷积层,用于将合并后的特征转换为输出通道 c2 ,使用1x1的卷积核和步长1。
self.cv7 = Conv(2 * c_, c2, 1, 1)
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 、 self.cv3 和 self.cv4 进行卷积操作,得到特征 x1 。
x1 = self.cv4(self.cv3(self.cv1(x)))
# 然后, x1 通过 self.cv5 和 self.cv6 进行进一步的处理,并且与通过 self.m 中的每个最大池化层处理后的 x1 进行拼接,得到特征 y1 。
y1 = self.cv6(self.cv5(torch.cat([x1] + [m(x1) for m in self.m], 1)))
# 同时, x 也通过 self.cv2 进行处理,得到特征 y2 。
y2 = self.cv2(x)
# 最后, y1 和 y2 在通道维度(维度1)上进行拼接,并通过 self.cv7 进行最终的卷积操作,得到最终的输出。
return self.cv7(torch.cat((y1, y2), dim=1))
# CSP和SPP的特点 :
# CSP : Cross Stage Partial Networks 通过将网络分成两个并行的路径,每个路径处理不同的特征,然后将它们合并,这样可以减少计算量并提高特征的表达能力。
# SPP : 空间金字塔池化层可以捕获不同尺度的特征,增强模型对不同尺度物体的识别能力。
# 在代码中, SPPCSPC 类结合了 CSP 和 SPP 的特点,通过并行处理和多尺度特征捕获,提高了网络的性能和效率。这种方法在保持模型性能的同时,增强了模型对不同尺度特征的捕获能力,适合用于深度学习模型中的特征提取阶段。
# 这段代码定义了一个名为 GhostSPPCSPC 的类,它是 SPPCSPC 类的子类。 GhostSPPCSPC 类继承了 SPPCSPC 类的所有属性和方法,并将其内部的卷积层替换为 GhostConv 层,这是一种GhostNet架构中提出的高效卷积结构。
class GhostSPPCSPC(SPPCSPC):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
# 4.shortcut :是否使用快捷连接,默认为False。
# 5.g :分组卷积的组数,默认为1。
# 6.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)):
super().__init__(c1, c2, n, shortcut, g, e, k)
c_ = int(2 * c2 * e) # hidden channels
self.cv1 = GhostConv(c1, c_, 1, 1)
self.cv2 = GhostConv(c1, c_, 1, 1)
self.cv3 = GhostConv(c_, c_, 3, 1)
self.cv4 = GhostConv(c_, c_, 1, 1)
self.cv5 = GhostConv(4 * c_, c_, 1, 1)
self.cv6 = GhostConv(c_, c_, 3, 1)
self.cv7 = GhostConv(2 * c_, c2, 1, 1)
# GhostConv 的特点 :
# GhostConv 是一种设计用来减少卷积层参数数量和计算量的卷积结构。它通过生成辅助特征图(cheap auxiliary channels)来增加网络的感受野,而不显著增加参数数量和计算量。
# 总结 :
# GhostSPPCSPC 类结合了 SPPCSPC 的网络结构和 GhostConv 的高效卷积机制,旨在提高网络的效率和性能。通过在CSP和SPP结构中使用 GhostConv , GhostSPPCSPC 可以在保持或提升性能的同时,减少模型的参数数量和计算量,适合用于需要高效率的深度学习模型中。
# 这段代码定义了一个名为 GhostStem 的类,它是 Stem 类的子类。 GhostStem 类继承了 Stem 类的所有属性和方法,并将其内部的卷积层替换为 GhostConv 层,这是一种GhostNet架构中提出的高效卷积结构。
class GhostStem(Stem):
# Stem
# 5.p :填充,默认为None。
# 6.g :分组卷积的组数,默认为1。
# 7.act :是否应用激活函数,默认为True。
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
super().__init__(c1, c2, k, s, p, g, act)
c_ = int(c2/2) # hidden channels
self.cv1 = GhostConv(c1, c_, 3, 2)
self.cv2 = GhostConv(c_, c_, 1, 1)
self.cv3 = GhostConv(c_, c_, 3, 2)
self.cv4 = GhostConv(2 * c_, c2, 1, 1)
# GhostStem 类结合了 Stem 的网络结构和 GhostConv 的高效卷积机制,旨在提高网络的效率和性能。
# 通过在Stem结构中使用 GhostConv , GhostStem 可以在保持或提升性能的同时,减少模型的参数数量和计算量,适合用于需要高效率的深度学习模型的初始层。
# 这种结构特别适用于图像分类和其他视觉任务的模型,因为它可以在网络的早期阶段有效地提取特征。
# 这段代码定义了一个名为 BottleneckCSPA 的类,它是PyTorch中的一个模块( nn.Module ),用于实现Cross Stage Partial Networks (CSP) 结合Darknet的Bottleneck结构。
# CSP是一种网络结构,旨在提高网络的效率和性能,而Bottleneck结构是一种常见的网络设计模式,用于减少网络的参数数量和计算量,同时保持或提升网络的性能。
class BottleneckCSPA(nn.Module):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
# 1.c1 :输入通道数( ch_in )。
# 2.c2 :输出通道数( ch_out )。
# 3.n :Bottleneck结构重复的次数。
# 4.shortcut :是否使用快捷连接(shortcut connection),默认为True。
# 5.g :分组卷积的组数,默认为1。
# 6.e :扩张率(expansion ratio),用于计算隐藏通道数,默认为0.5。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(BottleneckCSPA, self).__init__()
c_ = int(c2 * e) # hidden channels
# 第一个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,用于将输入通道 c1 转换为隐藏通道 c_ ,使用1x1的卷积核和步长1。
self.cv2 = Conv(c1, c_, 1, 1)
# 第三个卷积层,用于将两个分支的输出合并,并转换为输出通道 c2 ,使用1x1的卷积核和步长1。
self.cv3 = Conv(2 * c_, c2, 1, 1)
# 一个序列模块,包含 n 个Bottleneck结构,每个Bottleneck结构都是一个 Conv 层的序列,用于进一步处理特征。
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作,得到特征 y1 。
# 然后, y1 通过 self.m 中的 Bottleneck 结构进行进一步的处理。
y1 = self.m(self.cv1(x))
# 同时, x 也通过 self.cv2 进行卷积操作,得到特征 y2 。
y2 = self.cv2(x)
# 最后, y1 和 y2 在通道维度(维度1)上进行拼接,并通过 self.cv3 进行最终的卷积操作,得到最终的输出。
return self.cv3(torch.cat((y1, y2), dim=1))
# CSP 和 Bottleneck 的特点 :
# CSP :Cross Stage Partial Networks 通过将网络分成两个并行的路径,每个路径处理不同的特征,然后将它们合并,这样可以减少计算量并提高特征的表达能力。
# Bottleneck :Bottleneck 结构通过使用1x1卷积减少通道数,然后通过3x3卷积提取特征,最后再通过1x1卷积增加通道数,有效地捕获特征并减少计算量。
# 在代码中, BottleneckCSPA 类结合了 CSP 和 Bottleneck 的特点,通过并行处理和特征提取,提高了网络的性能和效率。这种方法在保持模型性能的同时,减少了模型的参数数量和计算量,适合用于深度学习模型中的中间层。
# 这段代码定义了一个名为 BottleneckCSPB 的类,它是PyTorch中的一个模块( nn.Module ),用于实现Cross Stage Partial Networks (CSP) 结合 Darknet 的 Bottleneck 结构。
# 这个类是对CSP结构的一个变种,其中使用了 Bottleneck 模块来构建网络的深度,同时保持计算效率。
class BottleneckCSPB(nn.Module):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(BottleneckCSPB, self).__init__()
c_ = int(c2) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
# 在前向传播方法中,输入数据 x 首先通过 self.cv1 进行卷积操作,得到特征 x1 。然后, x1 通过 self.m 中的 Bottleneck 结构进行进一步的处理,得到特征 y1 。
# 同时, x1 也通过 self.cv2 进行卷积操作,得到特征 y2 。最后, y1 和 y2 在通道维度(维度1)上进行拼接,并通过 self.cv3 进行最终的卷积操作,得到最终的输出。
x1 = self.cv1(x)
y1 = self.m(x1)
y2 = self.cv2(x1)
return self.cv3(torch.cat((y1, y2), dim=1))
# BottleneckCSPC 类是一个神经网络模块,它实现了Cross Stage Partial Networks (CSP) 的变体,这种结构在提高网络性能的同时减少了计算量。
class BottleneckCSPC(nn.Module):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.n : Bottleneck重复的次数,默认为1。
# 4.shortcut : 是否使用快捷连接,默认为True。
# 5.g : 分组卷积的组数,默认为1。
# 6.e : 扩张系数,默认为0.5。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(BottleneckCSPC, self).__init__()
c_ = int(c2 * e) # hidden channels
# 第一个卷积层,将输入 c1 通道转换为 c_ 通道。
self.cv1 = Conv(c1, c_, 1, 1)
# 第二个卷积层,同样将输入 c1 通道转换为 c_ 通道。
self.cv2 = Conv(c1, c_, 1, 1)
# 第三个卷积层,将 c_ 通道转换为 c_ 通道。
self.cv3 = Conv(c_, c_, 1, 1)
# 第四个卷积层,将 2 * c_ 通道转换为 c2 通道。
self.cv4 = Conv(2 * c_, c2, 1, 1)
# 一个序列模块,包含 n 个 Bottleneck 模块,每个模块的通道数为 c_ 。
self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# 定义了模块的前向传播过程。
def forward(self, x):
# 首先,输入 x 通过 cv1 卷积层,然后通过 m 序列模块(包含多个Bottleneck),最后通过 cv3 卷积层。
y1 = self.cv3(self.m(self.cv1(x)))
# 输入 x 通过 cv2 卷积层。
y2 = self.cv2(x)
# 最后, y1 和 y2 在通道维度上进行拼接,并通过 cv4 卷积层输出最终结果。
return self.cv4(torch.cat((y1, y2), dim=1))
# ResCSPA 类是继承自 BottleneckCSPA 的一个子类,它同样实现了Cross Stage Partial Networks (CSP) 的概念,但在这个变体中,它使用了残差连接(Residual Connections)来进一步增强网络的性能和训练稳定性。
class ResCSPA(BottleneckCSPA):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
# 一个序列模块,包含 n 个 Res 模块,每个模块的通道数为 c_ 。这里的 Res 模块是指残差模块,它将输入和输出通过一个或多个卷积层连接起来,以实现残差连接。
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# 总结 :
# ResCSPA 类通过在 BottleneckCSPA 的基础上引入残差连接,旨在提高网络的训练效率和性能。残差连接可以帮助网络学习更深层次的特征表示,同时避免梯度消失问题,这对于深层网络尤其重要。
# 这种结构在计算机视觉任务中,特别是在目标检测和图像分类中,已经被证明是非常有效的。
# ResCSPB 类是一个继承自 BottleneckCSPB 的PyTorch模块,它结合了残差连接(Residual Connections)和Cross Stage Partial Networks (CSP) 的概念。这种结构旨在提高网络的性能和训练稳定性,同时减少计算量。
class ResCSPB(BottleneckCSPB):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
# 计算隐藏层通道数,这是基于输出通道数 c2 和扩张系数 e 计算得出的。
c_ = int(c2) # hidden channels
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# 总结 :
# ResCSPB 类通过在 BottleneckCSPB 的基础上引入残差连接,旨在提高网络的训练效率和性能。残差连接可以帮助网络学习更深层次的特征表示,同时避免梯度消失问题,这对于深层网络尤其重要。
# 这种结构的引入,使得网络可以更有效地处理复杂的视觉任务,同时保持较高的计算效率。
# ResCSPC 类是继承自 BottleneckCSPC 的一个PyTorch模块,它结合了残差连接(Residual Connections)和Cross Stage Partial Networks (CSP) 的概念。这种结构旨在提高网络的性能和训练稳定性,同时减少计算量。
class ResCSPC(BottleneckCSPC):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# 这种结构的引入,使得网络可以更有效地处理复杂的视觉任务,同时保持较高的计算效率。通过结合CSP和残差连接, ResCSPC 类提供了一种强大的网络构建块,适用于需要高性能和高效率的深度学习应用。
class ResXCSPA(ResCSPA):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
# 5.g : 分组卷积的组数,默认为32。
def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# 总结 :
# ResXCSPA 类通过在 ResCSPA 的基础上进一步调整残差连接和扩张系数,旨在提高网络的训练效率和性能。这种结构在计算机视觉任务中,特别是在目标检测和图像分类中,已经被证明是非常有效的。
# 通过调整扩张系数和分组卷积的组数, ResXCSPA 类提供了更多的灵活性和性能优化空间。
class ResXCSPB(ResCSPB):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2) # hidden channels
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
class ResXCSPC(ResCSPC):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[Res(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# GhostCSPA 类是继承自 BottleneckCSPA 的一个PyTorch模块,它结合了Ghost模块和Cross Stage Partial Networks (CSP) 的概念。Ghost模块是一种用于减少网络计算复杂度的技术,它通过创建原始特征的廉价副本并应用不同的卷积核来实现。
class GhostCSPA(BottleneckCSPA):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])
class GhostCSPB(BottleneckCSPB):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2) # hidden channels
# 一个序列模块,包含 n 个 Ghost 模块,每个模块的通道数为 c_ 。这里的 Ghost 模块是指Ghost模块,它通过创建特征的副本并应用不同的卷积核来减少计算量。
self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])
# 总结 :
# GhostCSPA 类通过在 BottleneckCSPA 的基础上引入Ghost模块,旨在减少网络的计算复杂度,同时保持或提高网络的性能。Ghost模块通过创建特征的廉价副本并应用不同的卷积核,可以在不显著增加计算量的情况下增加网络的表达能力。
# 这种结构的引入,使得网络可以更有效地处理复杂的视觉任务,同时保持较高的计算效率。通过结合CSP和Ghost模块, GhostCSPA 类提供了一种强大的网络构建块,适用于需要高性能和高效率的深度学习应用。
class GhostCSPC(BottleneckCSPC):
# CSP https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[Ghost(c_, c_) for _ in range(n)])
##### end of cspnet #####
##### yolor #####
# ImplicitA 类是一个PyTorch模块,它实现了一个隐式层(Implicit Layer),这种层可以在不增加额外计算负担的情况下,为网络提供额外的学习能力。
class ImplicitA(nn.Module):
# 这是类的构造函数,用于初始化模块。
# 1.channel : 通道数,即输入和输出的特征图的通道数。
# 2.mean : 初始化参数的均值,默认为0。
# 3.std : 初始化参数的标准差,默认为0.02。
def __init__(self, channel, mean=0., std=.02):
super(ImplicitA, self).__init__()
# 存储输入和输出的特征图的通道数。
self.channel = channel
# 存储初始化参数的均值。
self.mean = mean
# 存储初始化参数的标准差。
self.std = std
# 一个可学习的参数,形状为 (1, channel, 1, 1) ,即一个在空间维度上为1x1的卷积核,这个参数会在训练过程中被优化。
self.implicit = nn.Parameter(torch.zeros(1, channel, 1, 1))
# torch.nn.init.normal_(tensor, mean=0., std=1.)
# nn.init.normal_ 是 PyTorch 中的一个函数,用于将正态分布(高斯分布)的值初始化到张量中。这个函数会就地修改张量的值,也就是说,它会直接改变传入的张量,而不是返回一个新的张量。
# 参数说明 :
# tensor : 要被初始化的张量。
# mean : 正态分布的均值(μ),默认为 0。
# std : 正态分布的标准差(σ),默认为 1。
# 这个函数会对 tensor 中的每个元素从均值为 mean 、标准差为 std 的正态分布中随机采样一个值,并赋值给对应的元素。
# 这种初始化方法常用于神经网络的权重初始化,因为它可以帮助打破对称性,并且有助于神经网络的梯度传播和收敛。
# 使用正态分布初始化 self.implicit 参数,均值为 mean ,标准差为 std 。
nn.init.normal_(self.implicit, mean=self.mean, std=self.std)
def forward(self, x):
# 输入 x 与 self.implicit 参数相加后输出。这种操作等效于在输入特征图上应用了一个1x1的卷积,但是这个卷积核是可学习的,并且在训练过程中会被优化。
return self.implicit + x
# ImplicitA 类通过引入一个可学习的隐式参数 self.implicit ,为网络提供了一种轻量级的调整输入特征的能力。这种隐式层不会增加额外的卷积计算,但是可以提供额外的自由度来优化网络的性能。
# 这种技术可以用于各种深度学习任务中,特别是在需要微调网络性能但又不希望显著增加计算负担的情况下。
# 注意事项 :确保输入 x 的通道数与 self.channel 相匹配,否则在执行加法操作时会出现维度不匹配的错误。
class ImplicitM(nn.Module):
def __init__(self, channel, mean=0., std=.02):
super(ImplicitM, self).__init__()
self.channel = channel
self.mean = mean
self.std = std
self.implicit = nn.Parameter(torch.ones(1, channel, 1, 1))
nn.init.normal_(self.implicit, mean=self.mean, std=self.std)
def forward(self, x):
# 输入 x 与 self.implicit 参数逐元素相乘后输出。这种操作等效于在输入特征图上应用了一个1x1的逐元素乘法操作,但是这个乘法因子是可学习的,并且在训练过程中会被优化。
return self.implicit * x
# 初始化 self.implicit 时使用了 torch.ones ,这意味着初始乘法因子为1,但随后通过 nn.init.normal_ 进行了正态分布的随机初始化,这可能会导致初始乘法因子不再是1。
# 这种隐式乘法层的设计允许模型在训练过程中学习如何最好地缩放输入特征,这可能有助于改善模型的表示能力。
##### end of yolor #####
##### repvgg #####
# 在YOLOv7中的 RepConv 类是一个可部署的卷积模块,它基于RepVGG中提出的概念。RepVGG是一种卷积神经网络架构,它在推理时具有VGG风格的结构,而在训练时则具有多分支拓扑结构。
# 这种结构的解耦是通过结构重参数化技术实现的,使得模型在训练和推理时可以有不同的架构。
class RepConv(nn.Module):
# Represented convolution
# https://arxiv.org/abs/2101.03697
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.k : 卷积核大小,默认为3。
# 4.s : 步长,默认为1。
# 5.p : 填充,默认为None,如果为None,则会自动计算以保持特征图尺寸。
# 6.g : 组数,默认为1,用于分组卷积。
# 7.act : 是否包含激活函数,默认为True。
# 8.deploy : 是否为部署模式,默认为False。在部署模式下,会将训练时的多分支结构转换为单一的卷积结构,以便于推理。
def __init__(self, c1, c2, k=3, s=1, p=None, g=1, act=True, deploy=False):
super(RepConv, self).__init__()
# 将传入的 deploy 参数值赋给类的 deploy 属性。这个属性决定了模块是否处于部署模式。
self.deploy = deploy
# 将传入的 g 参数值赋给类的 groups 属性。这个属性表示分组卷积的组数。
self.groups = g
# 将传入的 c1 参数值赋给类的 in_channels 属性。这个属性表示输入通道数。
self.in_channels = c1
# 将传入的 c2 参数值赋给类的 out_channels 属性。这个属性表示输出通道数。
self.out_channels = c2
# 断言卷积核大小 k 必须为3。这是因为RepConv是针对3x3卷积设计的。
assert k == 3
# 断言自动计算的填充 autopad(k, p) 必须为1。 autopad 函数的作用是根据卷积核大小 k 和步长 s 自动计算所需的填充,以保持特征图的尺寸。这里要求自动计算的填充为1,意味着卷积操作不会改变特征图的宽度和高度。
assert autopad(k, p) == 1
# 计算用于1x1卷积的填充。由于1x1卷积的卷积核大小是1,所以其填充应该是自动计算填充减去卷积核大小的一半(向下取整)。
padding_11 = autopad(k, p) - k // 2
# 初始化激活函数。
# 如果 act 为True,则使用LeakyReLU激活函数,斜率为0.1,并且设置为inplace操作以节省内存。
# 如果 act 是 nn.Module 的实例,则直接使用传入的激活函数模块。
# 如果 act 为False,则使用Identity模块,即不使用激活函数。
self.act = nn.LeakyReLU(0.1, inplace=True) if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
if deploy:
# 部署模式(deploy=True)。
# 如果处于部署模式, RepConv 会创建一个标准的2D卷积层( nn.Conv2d ),这个卷积层将输入通道数为 c1 的特征图转换为输出通道数为 c2 的特征图,使用卷积核大小 k 和步长 s 。
# autopad(k, p) 用于自动计算填充, groups=g 表示分组卷积。 bias=True 表示卷积层包含偏置项。
self.rbr_reparam = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=True)
else:
# 非部署模式(deploy=False)。
# 在非部署模式下,如果输出通道数 c2 等于输入通道数 c1 且步长 s 为1,那么会创建一个批量归一化层( nn.BatchNorm2d ),用于可能的身份(残差)连接。
# 如果不满足这个条件,则 rbr_identity 为 None 。
self.rbr_identity = (nn.BatchNorm2d(num_features=c1) if c2 == c1 and s == 1 else None)
# 创建一个序列容器( nn.Sequential ),包含一个2D卷积层和一个批量归一化层。这个卷积层使用卷积核大小 k ,步长 s ,自动计算的填充,分组数 g ,并且不包含偏置项( bias=False )。
self.rbr_dense = nn.Sequential(
nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False),
nn.BatchNorm2d(num_features=c2),
)
# 创建另一个序列容器,包含一个1x1的2D卷积层和一个批量归一化层。这个1x1卷积层使用步长 s 和之前计算的 padding_11 作为填充,分组数 g ,同样不包含偏置项。
self.rbr_1x1 = nn.Sequential(
nn.Conv2d( c1, c2, 1, s, padding_11, groups=g, bias=False),
nn.BatchNorm2d(num_features=c2),
)
# forward 方法决定了当给定输入数据时,模块如何处理这些数据。以下是对这段代码的详细解释。
def forward(self, inputs):
# 这个条件检查 RepConv 实例是否具有 rbr_reparam 属性。如果存在,说明模块处于部署模式。
if hasattr(self, "rbr_reparam"):
# 在部署模式下,直接使用重参数化的卷积层 rbr_reparam 处理输入数据,然后应用激活函数 self.act ,并将结果返回。
return self.act(self.rbr_reparam(inputs))
# 这个条件检查是否需要计算身份(残差)连接的输出。如果 self.rbr_identity 为 None ,则表示不需要身份连接。
if self.rbr_identity is None:
# 如果没有身份连接,则 id_out 设置为0。
id_out = 0
else:
id_out = self.rbr_identity(inputs)
# 在非部署模式下,使用 self.rbr_dense 和 self.rbr_1x1 分别处理输入数据,然后将这两个分支的输出与身份连接的输出( id_out )相加,最后应用激活函数 self.act 并将最终结果返回。
return self.act(self.rbr_dense(inputs) + self.rbr_1x1(inputs) + id_out)
# 定义了 RepConv 类中的 get_equivalent_kernel_bias 方法,该方法用于在部署模式下将训练时的多分支结构转换为一个等效的单一卷积结构。
def get_equivalent_kernel_bias(self):
# def _fuse_bn_tensor(self, branch):
# -> 用于将批量归一化(Batch Normalization, BN)层的参数融合到卷积层的权重中。这通常在将训练时的模型转换为部署模型时进行,以简化模型结构并提高推理效率。
# -> 返回融合后的卷积核( kernel * t )和偏置( beta - running_mean * gamma / std )。
# -> return kernel * t, beta - running_mean * gamma / std
# 对 rbr_dense 分支(通常是3x3卷积后接批量归一化)调用 _fuse_bn_tensor 方法,得到融合后的3x3卷积核和偏置。
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
# 对 rbr_1x1 分支(通常是1x1卷积后接批量归一化)调用 _fuse_bn_tensor 方法,得到融合后的1x1卷积核和偏置。
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
# 对 rbr_identity 分支(可能是用于残差连接的批量归一化层)调用 _fuse_bn_tensor 方法,得到融合后的身份卷积核和偏置。
kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
# 调用 _pad_1x1_to_3x3_tensor 方法,将1x1卷积核填充为3x3卷积核。这是必要的,因为我们需要将所有分支的卷积核合并为一个统一的3x3卷积核。
# 合并卷积核和偏置,将所有分支的卷积核和偏置相加,得到最终的等效卷积核和偏置。
return (
# def _pad_1x1_to_3x3_tensor(self, kernel1x1):
# -> 用于将 1x1 卷积层的权重(kernel)填充(pad)成 3x3 的形状。这在结构重参数化过程中非常有用,特别是在将训练时的多分支结构转换为部署时的单一卷积结构时。
# -> 如果 kernel1x1 不为 None ,则使用 PyTorch 的 pad 函数对其进行填充。 pad 函数的第二个参数 [1, 1, 1, 1] 指定了在每个维度上需要添加的填充量。
# -> return nn.functional.pad(kernel1x1, [1, 1, 1, 1])
kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid,
bias3x3 + bias1x1 + biasid,
)
# 这个方法的目的是将训练时的复杂结构(可能包括多个分支和批量归一化层)转换为一个简单的3x3卷积层,以便在部署时提高推理速度。
# 通过这种方式, RepConv 可以在部署模式下有效地模拟训练时的行为,同时保持模型的简化和高效。这种方法特别适用于在移动设备或嵌入式设备上部署模型,这些设备对计算资源和功耗有严格的限制。
# 这段代码定义了一个名为 _pad_1x1_to_3x3_tensor 的方法,它是 RepConv 类的一个辅助函数,用于将 1x1 卷积层的权重(kernel)填充(pad)成 3x3 的形状。这在结构重参数化过程中非常有用,特别是在将训练时的多分支结构转换为部署时的单一卷积结构时。
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
# 这个条件检查传入的 1x1 卷积核 kernel1x1 是否为 None 。如果为 None ,则表示没有 1x1 卷积分支,因此不需要进行任何操作。
if kernel1x1 is None:
return 0
else:
# torch.nn.functional.pad(input, pad, mode='xxx', value=x)
# torch.nn.functional.pad 是 PyTorch 中用于对张量进行填充操作的函数。填充操作在处理图像、序列数据等任务时非常常见,它可以在张量的指定维度两端添加一定数量的元素,填充方式多样,包括零填充、常数填充、反射填充和边界填充等。
# 返回一个新的张量,对输入张量进行了指定方式的填充。
# input (Tensor) :输入的张量。
# pad (tuple) :指定每个维度填充的数目。格式为 (left, right, top, bottom, …)。
# mode (str, 可选) :填充模式,包括 constant (常数填充,默认)、 reflect (反射填充)、 replicate (边界填充)和 circular( 循环填充)。
# value (float, 可选) :常数填充模式下填充值,仅当 mode 为 ‘ constant ’ 时有效,其他模式时即使指定了值也会被忽略掉。
# 如果 kernel1x1 不为 None ,则使用 PyTorch 的 pad 函数对其进行填充。 pad 函数的第二个参数 [1, 1, 1, 1] 指定了在每个维度上需要添加的填充量。
# 具体来说,这将在卷积核的上下左右四个方向上各添加一个单位的填充,从而将原始的 1x1 卷积核转换为 3x3 的形状。
return nn.functional.pad(kernel1x1, [1, 1, 1, 1])
# 这段代码定义了一个名为 _fuse_bn_tensor 的方法,它是 RepConv 类的一个辅助函数,用于将批量归一化(Batch Normalization, BN)层的参数融合到卷积层的权重中。这通常在将训练时的模型转换为部署模型时进行,以简化模型结构并提高推理效率。
def _fuse_bn_tensor(self, branch):
# 这个条件检查传入的分支 branch 是否为 None 。如果为 None ,则表示没有这个分支,因此返回两个零值(代表 融合后的 卷积核 和 偏置 )。
if branch is None:
return 0, 0
# 这个条件检查传入的分支是否是一个 nn.Sequential 序列容器。如果是,那么可以假设这个序列包含了 一个卷积层 后接 一个批量归一化层 。
if isinstance(branch, nn.Sequential):
# 从序列的第一个元素(卷积层)中提取权重。
kernel = branch[0].weight
# 提取批量归一化层的运行均值。
running_mean = branch[1].running_mean
# 提取批量归一化层的运行方差。
running_var = branch[1].running_var
# 提取批量归一化层的缩放因子(也称为权重)。
gamma = branch[1].weight
# 提取批量归一化层的偏移(也称为偏置)。
beta = branch[1].bias
# 提取批量归一化层的 epsilon 值,这是一个小的常数,用于防止除以零。
eps = branch[1].eps
# 它处理了另一种情况,即当 branch 是一个 nn.BatchNorm2d 实例时。这种情况下, branch 代表的是一个单独的批量归一化层,而不是一个包含卷积层和批量归一化层的 nn.Sequential 序列。
else:
# 确保 branch 确实是一个 nn.BatchNorm2d 实例。
assert isinstance(branch, nn.BatchNorm2d)
# 检查是否已经为这个实例初始化了 id_tensor 。如果没有,那么进行初始化。
if not hasattr(self, "id_tensor"):
# 计算每个组的输入维度。
input_dim = self.in_channels // self.groups
# 创建一个全零的数组,形状为 (self.in_channels, input_dim, 3, 3) ,用于存储身份张量。
kernel_value = np.zeros(
(self.in_channels, input_dim, 3, 3), dtype=np.float32
)
# 遍历输入通道,为每个通道在3x3卷积核中设置一个1,位置为 (i % input_dim, 1, 1) 。
# 在每个通道对应的3x3卷积核的中心位置设置为1。这样,每个输入通道只会映射到输出通道的相应位置,实现 身份映射 。
for i in range(self.in_channels):
kernel_value[i, i % input_dim, 1, 1] = 1
# 将numpy数组 kernel_value 转换为PyTorch张量,并将其移动到批量归一化层权重所在的设备(例如CPU或GPU)。这样, id_tensor 就可以在PyTorch的计算图中使用。
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
# 提取 身份映射 。
kernel = self.id_tensor
# 提取运行均值。
running_mean = branch.running_mean
# 提取运行方差。
running_var = branch.running_var
# 提取缩放因子(权重)。
gamma = branch.weight
# 提取偏移(偏置)。
beta = branch.bias
# 提取 epsilon 值。
eps = branch.eps
# 计算标准差。
std = (running_var + eps).sqrt()
# 计算缩放因子 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 类中的 repvgg_convert 方法,该方法用于将模型从训练时的多分支结构转换为部署时的单一卷积结构,并返回转换后的卷积核和偏置。
def repvgg_convert(self):
# def get_equivalent_kernel_bias(self):
# -> 该方法用于在部署模式下将训练时的多分支结构转换为一个等效的单一卷积结构。合并卷积核和偏置,将所有分支的卷积核和偏置相加,得到最终的等效卷积核和偏置。
# -> return (kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid,)
# 调用 get_equivalent_kernel_bias 方法,该方法将融合所有分支(包括3x3卷积、1x1卷积和可能的身份连接)的参数,返回一个等效的卷积核 kernel 和偏置 bias 。
kernel, bias = self.get_equivalent_kernel_bias()
# 返回 转换后的 卷积核 和 偏置 作为NumPy数组。
return (
# 将PyTorch张量 kernel 从当前计算图中分离( detach ),移动到CPU( cpu ),然后转换为NumPy数组( numpy() )。这是必要的,因为 NumPy 数组可以用于不同的深度学习框架或进行进一步的分析。
kernel.detach().cpu().numpy(),
# 同样地,将偏置 bias 从计算图中分离,移动到CPU,并转换为NumPy数组。
bias.detach().cpu().numpy(),
)
# 这个方法的目的是提供一个简单的方式,将 RepConv 模块中的复杂结构转换为一个等效的、可以直接用于推理的卷积层。这对于模型部署特别有用,因为它可以减少模型的复杂度,提高推理速度,同时保持模型的性能。
# 转换后的卷积核和偏置可以直接用于构建一个简化的卷积层,该层在结构上等同于原始的多分支结构,但在实际使用中更加高效。
# 这段代码定义了一个名为 fuse_conv_bn 的方法,它用于将一个卷积层( conv )和一个批量归一化层( bn )融合为一个单一的卷积层。这种融合操作通常在模型部署前进行,以简化模型结构并提高推理效率。
def fuse_conv_bn(self, conv, bn):
# 计算标准差,计算批量归一化层的运行方差加上 epsilon 值的 平方根 ,得到 标准差 。
std = (bn.running_var + bn.eps).sqrt()
# 计算偏置,根据批量归一化层的参数计算融合后的偏置项。
bias = bn.bias - bn.running_mean * bn.weight / std
# 计算缩放因子,计算批量归一化层的权重除以标准差,并将结果重塑为卷积核的形状,以便与卷积层的权重相乘。
t = (bn.weight / std).reshape(-1, 1, 1, 1)
# 计算融合后的权重,将卷积层的权重与缩放因子相乘,得到融合后的权重。
weights = conv.weight * t
bn = nn.Identity()
# 创建一个新的卷积层,创建一个新的 nn.Conv2d 层,其参数与原始卷积层相同,但包含融合后的权重和偏置,并且设置 bias 为 True 。
conv = nn.Conv2d(in_channels = conv.in_channels,
out_channels = conv.out_channels,
kernel_size = conv.kernel_size,
stride=conv.stride,
padding = conv.padding,
dilation = conv.dilation,
groups = conv.groups,
bias = True,
padding_mode = conv.padding_mode)
# 设置新的权重和偏置,将融合后的权重设置为新卷积层的权重。
conv.weight = torch.nn.Parameter(weights)
# 设置新的权重和偏置,将计算得到的偏置设置为新卷积层的偏置。
conv.bias = torch.nn.Parameter(bias)
# 返回新的卷积层,返回融合了批量归一化层参数的新卷积层。
return conv
# 这个方法的关键点在于,它将批量归一化层的效果直接编码到卷积层的权重和偏置中,从而在模型部署时可以省略批量归一化层,减少模型的计算负担。
# 这种融合操作在模型优化和部署时非常有用,特别是在需要将PyTorch模型转换为其他格式或在不支持批量归一化操作的硬件上运行时。
# 这段代码定义了 RepConv 类中的 fuse_repvgg_block 方法,该方法用于将 RepConv 模块中的所有分支(包括3x3卷积、1x1卷积和可能的身份连接)融合为一个单一的卷积结构,以便在部署模式下使用。
def fuse_repvgg_block(self):
if self.deploy:
# 检查是否处于部署模式,如果 RepConv 实例的 deploy 属性为 True ,表示模型已经在部署模式,不需要进行融合操作,直接返回。
return
# 打印融合信息,在控制台打印一条信息,表示正在执行 fuse_repvgg_block 方法。
print(f"RepConv.fuse_repvgg_block")
# def fuse_conv_bn(self, conv, bn):
# -> 它用于将一个卷积层( conv )和一个批量归一化层( bn )融合为一个单一的卷积层。这种融合操作通常在模型部署前进行,以简化模型结构并提高推理效率。
# -> 返回新的卷积层,返回融合了批量归一化层参数的新卷积层。
# -> return conv
# 融合3x3卷积和批量归一化,对 rbr_dense 分支(包含3x3卷积和批量归一化)调用 fuse_conv_bn 方法进行融合。 self.rbr_dense[0] 是卷积层, self.rbr_dense[1] 是批量归一化层。
self.rbr_dense = self.fuse_conv_bn(self.rbr_dense[0], self.rbr_dense[1])
# 融合1x1卷积和批量归一化,对 rbr_1x1 分支(包含1x1卷积和批量归一化)调用 fuse_conv_bn 方法进行融合。 self.rbr_1x1[0] 是卷积层, self.rbr_1x1[1] 是批量归一化层。
self.rbr_1x1 = self.fuse_conv_bn(self.rbr_1x1[0], self.rbr_1x1[1])
# 提取1x1卷积的偏置和权重。
# 获取融合后的1x1卷积层的偏置。
rbr_1x1_bias = self.rbr_1x1.bias
# 将1x1卷积核填充为3x3卷积核,以便与3x3卷积核进行合并。
weight_1x1_expanded = torch.nn.functional.pad(self.rbr_1x1.weight, [1, 1, 1, 1])
# 融合 self.rbr_identity。
# Fuse self.rbr_identity
# 检查 rbr_identity 类型,检查 rbr_identity 是否为 批量归一化层 或 同步批量归一化层 。
if (isinstance(self.rbr_identity, nn.BatchNorm2d) or isinstance(self.rbr_identity, nn.modules.batchnorm.SyncBatchNorm)):
# print(f"fuse: rbr_identity == BatchNorm2d or SyncBatchNorm")
# 创建1x1卷积层,创建一个1x1卷积层 identity_conv_1x1 ,其参数与 rbr_identity 对应的输入和输出通道数一致,不包含偏置项( bias=False )。
identity_conv_1x1 = nn.Conv2d(
in_channels=self.in_channels,
out_channels=self.out_channels,
kernel_size=1,
stride=1,
padding=0,
groups=self.groups,
bias=False)
# 初始化1x1卷积层权重。
# 将权重数据移动到与 rbr_1x1 权重相同的设备。
identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.to(self.rbr_1x1.weight.data.device)
# 移除权重数据中不必要的维度。
identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.squeeze().squeeze()
# print(f" identity_conv_1x1.weight = {identity_conv_1x1.weight.shape}")
# 将权重数据填充为0。
identity_conv_1x1.weight.data.fill_(0.0)
# 将权重数据的对角线设置为1,这样1x1卷积层实际上执行的是恒等变换。
identity_conv_1x1.weight.data.fill_diagonal_(1.0)
# 添加必要的维度,以使权重数据的形状与卷积核的形状一致。
identity_conv_1x1.weight.data = identity_conv_1x1.weight.data.unsqueeze(2).unsqueeze(3)
# print(f" identity_conv_1x1.weight = {identity_conv_1x1.weight.shape}")
# # def fuse_conv_bn(self, conv, bn):
# -> 它用于将一个卷积层( conv )和一个批量归一化层( bn )融合为一个单一的卷积层。这种融合操作通常在模型部署前进行,以简化模型结构并提高推理效率。
# -> 返回新的卷积层,返回融合了批量归一化层参数的新卷积层。
# -> return conv
# 融合1x1卷积和批量归一化。
# 调用 fuse_conv_bn 方法将1x1卷积层和批量归一化层融合。
identity_conv_1x1 = self.fuse_conv_bn(identity_conv_1x1, self.rbr_identity)
# 提取融合后的偏置和权重。
# 获取融合后的1x1卷积层的偏置。
bias_identity_expanded = identity_conv_1x1.bias
# 将融合后的1x1卷积核填充为3x3卷积核,以便与其他卷积核合并。
weight_identity_expanded = torch.nn.functional.pad(identity_conv_1x1.weight, [1, 1, 1, 1])
# 它处理了 rbr_identity 不是批量归一化层( nn.BatchNorm2d 或 nn.modules.batchnorm.SyncBatchNorm )的情况。
else:
# print(f"fuse: rbr_identity != BatchNorm2d, rbr_identity = {self.rbr_identity}")
# 处理非批量归一化层的 rbr_identity 。
# 如果 rbr_identity 不是批量归一化层,创建两个零参数, bias_identity_expanded 和 weight_identity_expanded ,它们的形状分别与 rbr_1x1 的偏置 rbr_1x1_bias 和 weight_1x1_expanded 相同。
# 这两个零参数将用于后续的合并操作,以确保所有分支的偏置和权重都能被正确地相加,即使 rbr_identity 不贡献任何非零值。
bias_identity_expanded = torch.nn.Parameter( torch.zeros_like(rbr_1x1_bias) )
weight_identity_expanded = torch.nn.Parameter( torch.zeros_like(weight_1x1_expanded) )
#print(f"self.rbr_1x1.weight = {self.rbr_1x1.weight.shape}, ")
#print(f"weight_1x1_expanded = {weight_1x1_expanded.shape}, ")
#print(f"self.rbr_dense.weight = {self.rbr_dense.weight.shape}, ")
# 合并所有分支的权重和偏置。
# 将3x3卷积核、1x1卷积核(已填充为3x3)和身份连接的卷积核(如果 rbr_identity 是批量归一化层,则为非零)相加,得到最终的 卷积核 。
self.rbr_dense.weight = torch.nn.Parameter(self.rbr_dense.weight + weight_1x1_expanded + weight_identity_expanded)
# 将3x3卷积层、1x1卷积层和身份连接的偏置相加,得到最终的 偏置 。
self.rbr_dense.bias = torch.nn.Parameter(self.rbr_dense.bias + rbr_1x1_bias + bias_identity_expanded)
# 设置部署模式。
# 将融合后的卷积层赋值给 rbr_reparam ,表示 RepConv 模块现在可以用一个单一的卷积层来表示。
self.rbr_reparam = self.rbr_dense
# 设置 deploy 标志为 True ,表示模块已转换为部署模式。
self.deploy = True
# 清理不再需要的分支。
# 如果 rbr_identity 、 rbr_1x1 或 rbr_dense 不为空,则删除这些分支,释放资源,并将其设置为 None 。这一步是为了确保在部署模式下,这些分支不再被使用,并且减少内存占用。
if self.rbr_identity is not None:
del self.rbr_identity
self.rbr_identity = None
if self.rbr_1x1 is not None:
del self.rbr_1x1
self.rbr_1x1 = None
if self.rbr_dense is not None:
del self.rbr_dense
self.rbr_dense = None
# 这个方法的目的是将 RepConv 模块中的复杂结构(可能包括多个分支和批量归一化层)转换为一个简单的3x3卷积层,以便在部署时提高推理速度。通过这种方式, RepConv 可以在部署模式下有效地模拟训练时的行为,同时保持模型的简化和高效。
# 定义了一个名为 RepBottleneck 的类,它继承自 Bottleneck 类,并在其中使用 RepConv 模块来构建一个标准的瓶颈结构。这种结构通常用于深度卷积神经网络中,以减少参数数量并增加网络深度。
class RepBottleneck(Bottleneck):
# Standard bottleneck
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.shortcut : 是否使用快捷连接(shortcut connection),默认为 True 。
# 4.g : 分组卷积的组数,默认为 1 。
# 5.e : 扩张系数,用于计算隐藏层通道数。
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
# 调用父类初始化方法。调用父类 Bottleneck 的初始化方法,传递初始化参数。
super().__init__(c1, c2, shortcut=True, g=1, e=0.5)
# 计算隐藏层通道数。计算隐藏层的通道数,这是基于输出通道数 c2 和扩张系数 e 的乘积。
c_ = int(c2 * e) # hidden channels
# 创建 RepConv 模块。
# 创建一个 RepConv 实例,用于构建瓶颈结构中的第二个卷积层。这个卷积层将隐藏层的通道数 c_ 转换为输出通道数 c2 ,使用3x3的卷积核和步长为1,分组数为 g 。
self.cv2 = RepConv(c_, c2, 3, 1, g=g)
# RepBottleneck 类通过使用 RepConv 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。
# 定义了一个名为 RepBottleneckCSPA 的类,它继承自 BottleneckCSPA 类,并在其中使用 RepBottleneck 模块来构建一个Cross Stage Partial Networks(CSPNet)中的瓶颈结构。CSPNet是一种高效的卷积神经网络架构,它通过跨阶段的部分连接来减少计算量。
class RepBottleneckCSPA(BottleneckCSPA):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.n : 重复的 RepBottleneck 模块的数量。
# 4.shortcut : 是否使用快捷连接(shortcut connection),默认为 True 。
# 5.g : 分组卷积的组数,默认为 1 。
# 6.e : 扩张系数,用于计算隐藏层通道数。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
# 调用父类初始化方法。调用父类 BottleneckCSPA 的初始化方法,传递初始化参数。
super().__init__(c1, c2, n, shortcut, g, e)
# 计算隐藏层通道数。
c_ = int(c2 * e) # hidden channels
# 创建 RepBottleneck 模块序列。
# 创建一个序列容器 self.m ,其中包含 n 个 RepBottleneck 实例。每个 RepBottleneck 实例将隐藏层的通道数 c_ 转换为相同的隐藏层通道数 c_ ,使用1.0作为扩张系数,这意味着每个 RepBottleneck 实际上没有扩张,即直接使用隐藏层通道数。
self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# RepBottleneckCSPA 类通过使用 RepBottleneck 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。
# 在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。通过这种方式, RepBottleneckCSPA 类实现了CSPNet的核心思想,即通过跨阶段的部分连接来提高网络的效率。
# 这段代码定义了一个名为 RepBottleneckCSPB 的类,它继承自 BottleneckCSPB 类,并在其中使用 RepBottleneck 模块来构建一个Cross Stage Partial Networks(CSPNet)中的瓶颈结构。CSPNet是一种高效的卷积神经网络架构,它通过跨阶段的部分连接来减少计算量。
class RepBottleneckCSPB(BottleneckCSPB):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2) # hidden channels
self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# 这段代码定义了一个名为 RepBottleneckCSPC 的类,它继承自 BottleneckCSPC 类,并在其中使用 RepBottleneck 模块来构建一个Cross Stage Partial Networks(CSPNet)中的瓶颈结构。这个结构是CSPNet的一个变体,它通过跨阶段的部分连接来减少计算量,同时保持特征提取的有效性。
class RepBottleneckCSPC(BottleneckCSPC):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[RepBottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# 这段代码定义了一个名为 RepRes 的类,它继承自 Res 类,并在其中使用 RepConv 模块来构建一个标准的残差网络(ResNet)中的瓶颈结构。这种结构是残差网络的一个变体,它通过使用 RepConv 模块来替代传统的卷积层,以提高网络的效率和性能。
class RepRes(Res):
# Standard bottleneck 标准瓶颈。
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.shortcut : 是否使用快捷连接(shortcut connection),默认为 True 。
# 4.g : 分组卷积的组数,默认为 1 。
# 5.e : 扩张系数,用于计算隐藏层通道数。
def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
# 调用父类初始化方法。调用父类 Res 的初始化方法,传递初始化参数。
super().__init__(c1, c2, shortcut, g, e)
# 计算隐藏层通道数。计算隐藏层的通道数,这是基于输出通道数 c2 和扩张系数 e 的乘积。
c_ = int(c2 * e) # hidden channels
# 创建 RepConv 模块。
# 创建一个 RepConv 实例,用于构建瓶颈结构中的第二个卷积层。这个卷积层将隐藏层的通道数 c_ 转换为相同的隐藏层通道数 c_ ,使用3x3的卷积核和步长为1,分组数为 g 。
self.cv2 = RepConv(c_, c_, 3, 1, g=g)
# RepRes 类通过使用 RepConv 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。
# 在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。
# 这段代码定义了一个名为 RepResCSPA 的类,它继承自 ResCSPA 类,并在其中使用 RepRes 模块来构建一个Cross Stage Partial Networks(CSPNet)中的残差瓶颈结构。CSPNet是一种高效的卷积神经网络架构,它通过跨阶段的部分连接来减少计算量。
class RepResCSPA(ResCSPA):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.n : 重复的 RepRes 模块的数量。
# 4.shortcut : 是否使用快捷连接(shortcut connection),默认为 True 。
# 5.g : 分组卷积的组数,默认为 1 。
# 6.e : 扩张系数,用于计算隐藏层通道数。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
# 调用父类初始化方法。
super().__init__(c1, c2, n, shortcut, g, e)
# 计算隐藏层通道数。
c_ = int(c2 * e) # hidden channels
# 创建 RepRes 模块序列。
# 创建一个序列容器 self.m ,其中包含 n 个 RepRes 实例。每个 RepRes 实例将隐藏层的通道数 c_ 转换为相同的隐藏层通道数 c_ ,使用0.5作为扩张系数。
self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# RepResCSPA 类通过使用 RepRes 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。
# 在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。通过这种方式, RepResCSPA 类实现了CSPNet的核心思想,即通过跨阶段的部分连接来提高网络的效率。
class RepResCSPB(ResCSPB):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2) # hidden channels
self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
class RepResCSPC(ResCSPC):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[RepRes(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# 这段代码定义了一个名为 RepResX 的类,它继承自 ResX 类,并在其中使用 RepConv 模块来构建一个标准瓶颈结构,这是一个变种的残差网络(ResNet)结构。
class RepResX(ResX):
# Standard bottleneck
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.shortcut : 是否使用快捷连接(shortcut connection),默认为 True 。
# 4.g : 分组卷积的组数,默认为 32 。
# 5.e : 扩张系数,用于计算隐藏层通道数。
def __init__(self, c1, c2, shortcut=True, g=32, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
# 调用父类初始化方法。
super().__init__(c1, c2, shortcut, g, e)
# 计算隐藏层通道数。
c_ = int(c2 * e) # hidden channels
# 创建 RepConv 模块。
# 创建一个 RepConv 实例,用于构建瓶颈结构中的第二个卷积层。这个卷积层将隐藏层的通道数 c_ 转换为相同的隐藏层通道数 c_ ,使用3x3的卷积核和步长为1,分组数为 g 。
self.cv2 = RepConv(c_, c_, 3, 1, g=g)
# RepResX 类通过使用 RepConv 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。
# 在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。
# 这个类的设计意图可能是为了提供一个更灵活的残差结构,其中 g 参数(分组卷积的组数)被设置为一个相对较大值(默认为32),这可能有助于在保持性能的同时减少参数数量和计算量。
# 扩张系数 e 用于调整隐藏层的通道数,从而可以调整网络的容量和计算复杂度。
# 这段代码定义了一个名为 RepResXCSPA 的类,它继承自 ResXCSPA 类,并在其中使用 RepResX 模块来构建一个Cross Stage Partial Networks(CSPNet)中的残差瓶颈结构。这个结构结合了残差网络(ResNet)的特点和CSPNet的跨阶段部分连接,以提高网络的效率和性能。
class RepResXCSPA(ResXCSPA):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
# 初始化参数。
def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
# 调用父类初始化方法。
super().__init__(c1, c2, n, shortcut, g, e)
# 计算隐藏层通道数。
c_ = int(c2 * e) # hidden channels
# 创建 RepResX 模块序列。
# 创建一个序列容器 self.m ,其中包含 n 个 RepResX 实例。每个 RepResX 实例将隐藏层的通道数 c_ 转换为相同的隐藏层通道数 c_ ,使用0.5作为扩张系数,分组数为 g 。
self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
# RepResXCSPA 类通过使用 RepResX 模块替代传统的卷积层,可以在保持性能的同时减少模型的复杂度和计算量。这种设计特别适合于需要高效率和高速度的深度学习应用场景。
# 在部署模式下, RepConv 可以将多个分支融合为一个单一的卷积层,进一步简化模型结构。
# 这个类的设计意图可能是为了提供一个更灵活的残差结构,其中 g 参数(分组卷积的组数)被设置为一个相对较大值(默认为32),这可能有助于在保持性能的同时减少参数数量和计算量。
# 扩张系数 e 用于调整隐藏层的通道数,从而可以调整网络的容量和计算复杂度。
class RepResXCSPB(ResXCSPB):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2) # hidden channels
self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
class RepResXCSPC(ResXCSPC):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=32, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*[RepResX(c_, c_, shortcut, g, e=0.5) for _ in range(n)])
##### end of repvgg #####
##### transformer #####
# 这段代码定义了一个名为 TransformerLayer 的类,它是一个基于 PyTorch 的 nn.Module 的神经网络模块,用于实现 Transformer 架构中的一个单独的层。这个类不包含 LayerNorm 层,因为在某些情况下,移除 LayerNorm 层可能会提高性能。
class TransformerLayer(nn.Module):
# Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance)
# 初始化方法 __init__ 。
# 1.c : 表示每个注意力头的维度(embed_dim)。
# 2.num_heads : 表示多头注意力机制中头的数量。
def __init__(self, c, num_heads):
super().__init__()
# 属性。
# 一个线性层,用于计算查询(query)向量。它的输入和输出维度都是 c ,没有偏置项(bias=False)。
self.q = nn.Linear(c, c, bias=False)
# 一个线性层,用于计算键(key)向量。同样,输入和输出维度都是 c ,没有偏置项。
self.k = nn.Linear(c, c, bias=False)
# 线性层,用于计算值(value)向量。输入和输出维度也是 c ,没有偏置项。
self.v = nn.Linear(c, c, bias=False)
# torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None)
# 功能:创建一个多头注意力模块。
# embed_dim :输入数据的维度,也就是向量的长度。
# num_heads :表示并行注意力的数量,也就是“头”的数量。
# dropout :表示注意力权重的丢弃概率,相当于生成注意力之后,再将注意力传入一层Dropout层,默认为0。
# bias :在做线性变换Linear时,是否添加偏置,默认True。
# add_bias_kv :kv做线性变换时是否加偏置,若键值维度与嵌入维度相同,则可以将add_bias_kv设为False,默认False。
# add_zero_attn :将一个全零注意力向量添加到最终的输出中(只影响形状,不改变数值),强制使输出张量的形状与输入张量相同。
# kdim :keys的特征数据维度,即向量长度,默认与embed_dim相等。
# vdim :values的特征数据维度,即向量长度,默认与embed_dim相等。
# batch_first :如果设为True,则输入、输出张量表示为(batch, seq, feature),否则张量表示为(seq, batch, feature),默认False。
# 多头注意力层,使用 nn.MultiheadAttention 实现。 embed_dim 参数设置为 c , num_heads 设置为 num_heads ,表示每个头处理的维度。
self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads)
# 两个线性层,用于实现 Transformer 层中的前馈网络(feed-forward network)。这两个层的输入和输出维度都是 c ,没有偏置项。
self.fc1 = nn.Linear(c, c, bias=False)
self.fc2 = nn.Linear(c, c, bias=False)
# 这个 TransformerLayer 类可以被用来构建 Vision Transformer(ViT)模型中的一个层,ViT 模型将图像分割成固定大小的补丁序列,然后将这些补丁的线性嵌入作为输入传递给 Transformer 编码器。
# 这个类实现了 Transformer 编码器的一个关键组成部分,包括多头自注意力机制和前馈网络。在实际使用时,可以通过堆叠多个这样的层来构建完整的 Transformer 编码器。
# 定义了 TransformerLayer 类的 forward 方法,它指定了当给定输入 x 时,这个 Transformer 层如何处理这些数据。
def forward(self, x):
# 自注意力机制。
# x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x :这一行执行多头自注意力操作。
# self.q(x) :计算输入 x 的查询(query)向量。
# self.k(x) :计算输入 x 的键(key)向量。
# self.v(x) :计算输入 x 的值(value)向量。
# self.ma(...) :将查询、键、值传递给 nn.MultiheadAttention 层,执行自注意力计算。 ma 方法返回的是一个元组,其中第一个元素是注意力输出,因此使用 [0] 来获取这个输出。
# + x :将自注意力的输出与原始输入 x 相加,实现残差连接。
x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x
# 前馈网络。
# x = self.fc2(self.fc1(x)) + x :这一行执行前馈网络操作。
# self.fc1(x) :第一个线性层处理输入 x 。
# self.fc2(...) :第二个线性层处理第一个线性层的输出。
# + x :将前馈网络的输出与自注意力的输出(也是原始输入 x )相加,再次实现残差连接。
x = self.fc2(self.fc1(x)) + x
# 返回结果。
# 返回最终的输出 x ,它包含了自注意力和前馈网络的计算结果,并且通过残差连接与原始输入相结合。
return x
# 这个方法的实现确保了 Transformer 层能够处理输入数据,并通过自注意力机制和前馈网络来更新这些数据。残差连接有助于避免在深层网络中出现的梯度消失问题,并且使得网络能够学习到更有效的特征表示。
# 这种结构是 Transformer 模型的核心,使其在处理序列数据时表现出色,包括自然语言处理和计算机视觉任务。
# 这段代码定义了一个名为 TransformerBlock 的类,它是一个基于 PyTorch 的 nn.Module 的神经网络模块,用于实现 Vision Transformer(ViT)中的一个关键组件。这个类结合了卷积层(如果需要的话)和多个 Transformer 层,以处理图像数据。
class TransformerBlock(nn.Module):
# Vision Transformer https://arxiv.org/abs/2010.11929
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.num_heads : 多头自注意力机制中头的数量。
# 4.num_layers : Transformer 层的数量。
def __init__(self, c1, c2, num_heads, num_layers):
# 父类初始化。调用父类 nn.Module 的初始化方法。
super().__init__()
# 卷积层。初始化一个卷积层变量,默认为 None 。
self.conv = None
# 如果输入通道数 c1 不等于输出通道数 c2 ,则创建一个卷积层。
if c1 != c2:
# 使用自定义的 Conv 类(或函数)创建一个卷积层,将输入通道数从 c1 转换为 c2 。
self.conv = Conv(c1, c2)
# 线性层。创建一个线性层(全连接层),用于学习位置嵌入。这个层的输入和输出维度都是 c2 。
self.linear = nn.Linear(c2, c2) # learnable position embedding
# Transformer 层序列。
# 创建一个序列容器,包含 num_layers 个 TransformerLayer 实例。每个 TransformerLayer 都使用 c2 作为特征维度和 num_heads 作为头数。
self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)])
# 属性。保存输出通道数 c2 作为实例属性。
self.c2 = c2
# 这个 TransformerBlock 类的设计允许它灵活地处理不同维度的输入和输出,并且可以通过堆叠多个 Transformer 层来构建深度模型。如果输入和输出通道数不同,它还会应用一个卷积层来进行维度转换。
# 这种设计使得 TransformerBlock 成为构建 Vision Transformer 模型的一个非常有用的构建块。
# 定义了 TransformerBlock 类的 forward 方法,它指定了当给定输入 x 时,这个 Transformer 块如何处理这些数据。
def forward(self, x):
# 卷积层处理。
# 检查是否定义了卷积层。
if self.conv is not None:
# 如果存在卷积层,则将输入 x 通过卷积层进行处理。
x = self.conv(x)
# 输入维度调整。
# 获取输入 x 的批次大小 b 、通道数(忽略)、宽度 w 和高度 h 。
b, _, w, h = x.shape
# 将输入 x 的空间维度(宽度和高度)展平。
p = x.flatten(2)
# 在展平的张量 p 前添加一个新的维度,为后续的转置操作做准备。
p = p.unsqueeze(0)
# 将张量 p 的第一维和最后一维进行转置,以适应位置嵌入的维度要求。
p = p.transpose(0, 3)
# 移除转置后多余的单一维度。
p = p.squeeze(3)
# 位置嵌入。
# 通过线性层(位置嵌入层)计算位置嵌入 e 。
e = self.linear(p)
# 将位置嵌入 e 添加到展平的输入 p 上。
x = p + e
# Transformer 层处理。将添加了位置嵌入的张量 x 通过 Transformer 层序列进行处理。
x = self.tr(x)
# 输出维度恢复。
# 在处理后的张量 x 后添加一个新的维度,为后续的转置操作做准备。
x = x.unsqueeze(3)
# 将张量 x 的第一维和最后一维进行转置,以恢复到原始的批次和通道维度。
x = x.transpose(0, 3)
# 将张量 x 重塑为原始的批次大小 b 、通道数 self.c2 、宽度 w 和高度 h 。
x = x.reshape(b, self.c2, w, h)
# 返回结果。返回处理后的输出 x ,它具有与输入相同的空间维度,但已经通过了 Transformer 块的处理。
return x
# 这个方法的实现确保了输入数据在通过 Transformer 块时,其空间结构被保留,并且能够接收和输出具有相同空间维度的特征图。这种设计使得 TransformerBlock 可以灵活地集成到各种视觉任务的网络架构中。
##### end of transformer #####
##### yolov5 #####
# 定义了一个名为 Focus 的类,它是一个 PyTorch nn.Module 的子类,用于将空间信息(宽高)聚焦到通道空间。这个类的设计灵感可能来自于 EfficientNet 中的 Focus 层,它通过重新排列和组合输入特征图的相邻像素来增加通道维度,同时减少空间维度。
class Focus(nn.Module):
# 将 宽高 信息聚焦到 通道 空间。
# Focus wh information into c-space
# 初始化参数。
# 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): # ch_in, ch_out, kernel, stride, padding, groups
# 父类初始化。调用父类 nn.Module 的初始化方法。
super(Focus, self).__init__()
# 卷积层。创建一个卷积层,将输入通道数从 c1 * 4 转换为 c2 。这里 c1 * 4 是因为 Focus 层会将输入特征图的相邻像素组合起来,从而增加通道数。
self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
# 注释掉的代码。注释掉了一个 Contract 层,这个层可能用于进一步降低空间维度,但在这里没有被使用。
# self.contract = Contract(gain=2)
# 前向传播。定义前向传播方法。
def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
# x[..., ::2, ::2] : 取出输入特征图 x 的每个通道的每隔一个像素(水平和垂直)。
# x[..., 1::2, ::2] : 取出输入特征图 x 的每个通道的每隔一个像素(水平),但偏移一个像素。
# x[..., ::2, 1::2] : 取出输入特征图 x 的每个通道的每隔一个像素(垂直),但偏移一个像素。
# x[..., 1::2, 1::2] : 取出输入特征图 x 的每隔一个像素(水平和垂直),偏移一个像素。
# torch.cat([...], 1) : 将上述四个张量沿着通道维度(维度 1)拼接起来,形成新的输入特征图。
# self.conv(...) : 将拼接后的特征图通过卷积层进行处理。
return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
# 注释掉的代码。注释掉了使用 Contract 层的前向传播代码。
# return self.conv(self.contract(x))
# Focus 类的设计目的是在不增加计算复杂度的情况下,通过重新排列和组合输入特征图的相邻像素来增加通道维度,同时减少空间维度。这种方法有助于模型捕获更多的上下文信息,并且可以有效地减少参数数量和计算量。
# 这段代码定义了一个名为 SPPF 的类,它是一个 PyTorch nn.Module 的子类,用于实现空间金字塔池化(Spatial Pyramid Pooling)的快速版本(SPPF)。这种层通常用于计算机视觉任务中,特别是在目标检测模型如 YOLOv5 中,用于提取多尺度的特征。
class SPPF(nn.Module):
# 空间金字塔池化 - Glenn Jocher 为 YOLOv5 提供的快速 (SPPF) 层。
# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
# 初始化参数。
# 1.c1 : 输入通道数。
# 2.c2 : 输出通道数。
# 3.k : 池化核大小,默认为 5,相当于传统的 SPP 中的 (5, 9, 13)。
def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
# 父类初始化。调用父类 nn.Module 的初始化方法。
super().__init__()
c_ = c1 // 2 # hidden channels
# 卷积层。
# 创建一个卷积层,将输入通道数从 c1 转换为 c_ ,其中 c_ 是 c1 的一半,卷积核大小为 1,步长为 1。
self.cv1 = Conv(c1, c_, 1, 1)
# 创建另一个卷积层,将 c_ * 4 通道转换为 c2 通道,卷积核大小为 1,步长为 1。
self.cv2 = Conv(c_ * 4, c2, 1, 1)
# 最大池化层。创建一个最大池化层,核大小为 k ,步长为 1,填充为 k // 2 ,以保持输出尺寸不变。
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
# 前向传播。 定义前向传播方法。
def forward(self, x):
# 将输入 x 通过第一个卷积层 cv1 。
x = self.cv1(x)
# 将 x 通过最大池化层 m 得到 y1 。
y1 = self.m(x)
# 将 y1 再次通过最大池化层 m 得到 y2 。
y2 = self.m(y1)
# 将 x 、 y1 、 y2 和 self.m(y2) 沿着通道维度(维度 1)拼接起来,然后通过第二个卷积层 cv2 进行处理,得到最终输出。
return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1))
# SPPF 类通过连续的池化和卷积操作,有效地捕获了不同尺度的特征,这对于目标检测等任务非常有帮助。这种设计可以减少模型的计算负担,同时保持或提高特征提取的能力。
# 这段代码定义了一个名为 Contract 的类,它是一个 PyTorch nn.Module 的子类,用于将图像的空间维度(宽和高)压缩到通道维度。这种操作通常用于减少特征图的空间分辨率,同时增加通道数,这在某些网络架构中有助于捕捉更抽象的特征表示。
class Contract(nn.Module):
# 将宽度-高度收缩到通道中,即 x(1,64,80,80) 到 x(1,256,40,40)
# Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
# 初始化参数。
# 1.gain : 压缩因子,默认为 2。这意味着空间维度将被压缩为原来的一半。
def __init__(self, gain=2):
# 父类初始化。调用父类 nn.Module 的初始化方法。
super().__init__()
# 属性。保存压缩因子。
self.gain = gain
# 前向传播。定义前向传播方法。
def forward(self, x):
# 获取输入 x 的批次大小 N 、通道数 C 、高度 H 和宽度 W 。
N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
# 获取压缩因子 s 。
s = self.gain
# 将输入 x 重塑为 N 批次、 C 通道、 H // s 行、 s 行高、 W // s 列、 s 列宽的形状。
x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)
# 重新排列 x 的维度,将行高和列宽的维度移动到通道维度,变为 N 批次、 s 行高、 s 列宽、 C 通道、 H // s 高、 W // s 宽的形状。
x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
# 最后,将 x 重塑为 N 批次、 C * s * s 通道、 H // s 高、 W // s 宽的形状。
return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)
# Contract 类通过这种重塑和重新排列操作,实现了将空间维度压缩到通道维度的目的。这种操作可以减少特征图的空间尺寸,同时增加通道数,有助于网络学习更高层次的特征表示。
# 这种类型的操作在某些网络架构中,如 EfficientNet 和 MobileNet 等,被用来减少计算量和参数数量,同时保持或提高性能。
# 这段代码定义了一个名为 Expand 的类,它是 PyTorch 的 nn.Module 的子类,用于将通道维度扩展到空间维度(宽度和高度)。这种操作通常用于增加特征图的空间分辨率,同时减少通道数,这在某些网络架构中有助于恢复空间信息或进行上采样。
class Expand(nn.Module):
# 将通道扩展为宽度-高度,即 x(1,64,80,80) 到 x(1,16,160,160)。
# Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
# 初始化参数。
# 1.gain : 扩展因子,默认为 2。这意味着空间维度将被扩展为原来的两倍。
def __init__(self, gain=2):
# 父类初始化。 调用父类 nn.Module 的初始化方法。
super().__init__()
# 属性。保存扩展因子。
self.gain = gain
# 前向传播。定义前向传播方法。
def forward(self, x):
# 获取输入 x 的批次大小 N 、通道数 C 、高度 H 和宽度 W 。
N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
# 获取扩展因子 s 。
s = self.gain
# 将输入 x 重塑为 N 批次、 s 行、 s 列、 C // s ** 2 通道、 H 高、 W 宽的形状。
x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80)
# 重新排列 x 的维度,将行和列的维度移动到通道维度之后,变为 N 批次、 C // s ** 2 通道、 H 高、 s 行、 W 宽、 s 列的形状。
x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
# 最后,将 x 重塑为 N 批次、 C // s ** 2 通道、 H * s 高、 W * s 宽的形状。
return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160)
# Expand 类通过这种重塑和重新排列操作,实现了将通道维度扩展到空间维度的目的。这种操作可以增加特征图的空间尺寸,同时减少通道数,有助于网络在需要恢复细节信息的任务中表现更好,例如图像分割、超分辨率等。
# 这种类型的操作在某些网络架构中,如 U-Net、SegNet 等,被用来增加特征图的分辨率,以便更好地保留空间信息。
# 这段代码定义了一个名为 NMS 的类,它是 PyTorch 的 nn.Module 的子类,用于实现非最大抑制(Non-Maximum Suppression, NMS)操作。NMS 是一种在目标检测任务中常用的技术,用于去除重叠的检测框,保留最有可能包含目标的框。
class NMS(nn.Module):
# 非最大抑制 (NMS) 模块。
# Non-Maximum Suppression (NMS) module
# 类属性。
# 置信度阈值,只有置信度高于这个值的检测框才会被考虑。
conf = 0.25 # confidence threshold 置信阈值。
# 交并比(IoU)阈值,用于确定两个检测框是否重叠过多。
iou = 0.45 # IoU threshold IoU阈值。
# 一个可选的列表,用于过滤特定类别的检测框。
classes = None # (optional list) filter by class (可选列表)按类别过滤。
# 初始化方法。初始化方法,不接收任何参数。
def __init__(self):
# 调用父类 nn.Module 的初始化方法。
super(NMS, self).__init__()
# 前向传播。定义前向传播方法,接收输入 x 。
def forward(self, x):
# def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, labels=()):
# -> 该函数实现了非最大抑制(NMS)算法,用于在目标检测任务中去除重叠的检测框。返回包含NMS处理结果的输出张量 output 。
# -> return output
# 调用 non_max_suppression 函数,该函数实现了 NMS 算法。
# 输入 x[0] 应该是一个包含检测框、置信度和类别概率的张量。 conf_thres 和 iou_thres 分别是类属性中定义的置信度阈值和 IoU 阈值, classes 是可选的过滤类别列表。
return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
# 这段代码定义了一个名为 autoShape 的类,它是一个用于目标检测模型的包装器(wrapper),旨在使模型能够接受不同格式的输入(如OpenCV、NumPy、PIL图像或PyTorch张量),并包括预处理、推理和非极大值抑制(NMS)。
# 类定义。autoShape 继承自 nn.Module ,这是PyTorch中所有神经网络模块的基类。
class autoShape(nn.Module):
# 用于传递 cv2/np/PIL/torch 输入的输入稳健模型包装器。包括预处理、推理和 NMS。
# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
# 类属性。
# conf :NMS的置信度阈值,用于过滤检测结果。
conf = 0.25 # NMS confidence threshold
# iou :NMS的交并比阈值,用于确定是否抑制重叠的边界框。
iou = 0.45 # NMS IoU threshold
# classes :一个可选的列表,用于过滤特定类别的检测结果。
classes = None # (optional list) filter by class
# 构造函数。
# 1.model :是一个预训练的目标检测模型。
def __init__(self, model):
# 调用父类的构造函数,这是Python中继承机制的一部分。
super(autoShape, self).__init__()
# model.eval()
# .eval() 是一个用于将模型设置为评估模式的方法。这意味着模型将不会进行梯度计算和反向传播,这对于模型的推理阶段至关重要。
# 为什么要使用 .eval() ?
# 1. 梯度计算:在训练过程中,模型需要计算梯度以更新权重。然而,在评估阶段,梯度计算是不必要的,这会增加计算负担并消耗更多的内存和时间。
# 2. 批归一化(Batch Normalization):.eval()影响模型中某些层的行为,尤其是那些涉及统计数据的层,如批归一化层。在训练模式下,批归一化层会计算小批量数据的均值和方差,而在评估模式下,它将使用训练阶段计算的移动平均值。
# 3. 丢弃(Dropout):在训练过程中,丢弃层有助于防止过拟合,通过随机关闭一些神经元来增加模型的泛化能力。但在评估模式下,丢弃层将不再关闭任何神经元,因为我们需要模型以最佳状态运行。
# 4. 权重衰减(Weight Decay):在训练过程中,权重衰减用于正则化模型,但在评估时不需要应用。
# 将传入的模型设置为评估模式。在PyTorch中, .eval() 方法用于将模型设置为评估模式,这会关闭dropout和batch normalization等仅在训练时使用的层。
self.model = model.eval()
# 这段代码定义了一个名为 autoshape 的方法,它是 autoShape 类的一个实例方法。这个方法的目的是检查 autoShape 功能是否已经被启用,并且如果已经启用,则打印一条消息并返回当前实例。
# autoshape 是一个实例方法,它不接受除了 self 之外的任何参数。
def autoshape(self):
# 当调用 autoshape 方法时,它会打印一条消息,表明 autoShape 功能已经被启用,因此不需要再次转换模型。
# 自动形状已启用,跳过...
print('autoShape already enabled, skipping... ') # model already converted to model.autoshape() 模型已转换为 model.autoshape()。
# 返回值。方法返回 self ,即当前类的实例。这允许方法链式调用,例如 model.autoshape() 可以直接调用而不需要接收返回值。
return self
# 使用场景 :
# 这个方法可能用于确保 autoShape 类只被应用一次。如果用户不小心多次调用了 autoshape 方法,这个方法会提醒用户并避免重复应用 autoShape 功能。
# 这个方法使用了PyTorch的 torch.no_grad() 装饰器,这意味着在推理过程中不会计算梯度,这有助于减少内存消耗并加速推理过程。这个装饰器告诉PyTorch在这个方法中不需要计算梯度,这对于模型推理来说是常见的做法。
@torch.no_grad()
# 方法定义。 forward 方法是PyTorch模型中的一个特殊方法,用于定义模型的前向传播逻辑。
# 1.imgs :参数是输入图像,可以是文件名、URI、OpenCV图像、PIL图像、NumPy数组或PyTorch张量。
# 2.size :参数是目标图像尺寸,默认为640(通常用于宽度为1280的图像)。
# 3.augment :参数用于数据增强。
# 4.profile :参数用于性能分析。
def forward(self, imgs, size=640, augment=False, profile=False):
# 从各种来源推断。对于高度=640、宽度=1280,RGB 图像示例输入为:
# 文件名:imgs = 'data/samples/zidane.jpg'
# URI:= 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
# OpenCV:= cv2.imread('image.jpg')[:,:,::-1] # HWC BGR 到 RGB x(640,1280,3)
# PIL:= Image.open('image.jpg') # HWC x(640,1280,3)
# numpy:= np.zeros((640,1280,3)) # HWC
# torch:= torch.zeros(16,3,320,640) # BCHW(缩放到大小=640,值为 0-1)
# multiple:= [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # 图像列表
# Inference from various sources. For height=640, width=1280, RGB images example inputs are:
# filename: imgs = 'data/samples/zidane.jpg'
# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3)
# PIL: = Image.open('image.jpg') # HWC x(640,1280,3)
# numpy: = np.zeros((640,1280,3)) # HWC
# torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values)
# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
# 输入处理。
# def time_synchronized(): -> 它用于获取一个与CUDA时间同步的准确时间戳。在调用 time.time() 之前同步CUDA设备,可以确保获取的时间戳是在GPU操作完成后的,这对于性能分析和时间测量是非常重要的。 -> return time.time()
# t 用于记录时间, time_synchronized() 是一个自定义函数,用于获取当前时间。
t = [time_synchronized()]
# p 获取模型的第一个参数,用于确定模型所在的设备(CPU或GPU)和数据类型。
p = next(self.model.parameters()) # for device and type
if isinstance(imgs, torch.Tensor): # torch
# torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)
# torch.cuda.amp.autocast 是 PyTorch 中用于自动混合精度(Automatic Mixed Precision, AMP)训练的上下文管理器。它允许模型在训练过程中动态选择使用单精度(FP32)或半精度(FP16)进行计算,以此来平衡计算速度、内存使用和数值精度。
# 参数说明 :
# enabled (bool): 是否启用自动混合精度。默认为 True 。
# dtype (torch.dtype): 指定自动转换的目标数据类型。如果指定,就以该类型为准;如果没有指定,则根据 device_type 为 "cuda" 时默认为 torch.float16 ,为 "cpu" 时默认为 torch.bfloat16 。
# cache_enabled (bool): 是否启用缓存。默认为 True 。当启用时,可以提高性能,但可能会增加一些内存开销。
# 工作原理 :
# 上下文管理器: autocast 可以作为上下文管理器使用,它会影响代码块内的运算,使其在 FP16 或 FP32 之间自动切换。 with autocast():
# 装饰器: autocast 也可以作为装饰器使用,应用于模型的 forward 方法或其他需要自动混合精度的函数。 @autocast()
# 如果输入 imgs 是一个PyTorch张量,那么使用 amp.autocast 来优化推理过程。 amp.autocast 是PyTorch中的自动混合精度工具,它可以在保持模型精度的同时减少内存消耗和加速推理。
with amp.autocast(enabled=p.device.type != 'cpu'):
# imgs.to(p.device) 将输入张量移动到模型所在的设备。
# imgs.type_as(p) 将输入张量转换为模型参数的数据类型。
# self.model(...) 调用模型的前向传播方法,传入处理后的图像和控制参数。
return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
# Pre-process
# 确定图像数量和列表。这行代码检查 imgs 是否是一个列表,如果是,则直接使用;如果不是,则将其放入一个列表中,并设置图像数量 n 为1。
n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
# 初始化变量。初始化三个列表,用于存储 原始图像尺寸 、 调整后的图像尺寸 和 文件名 。
shape0, shape1, files = [], [], [] # image and inference shapes, filenames
# 遍历图像列表。使用 enumerate 函数遍历图像列表 imgs , i 是索引, im 是当前图像。
for i, im in enumerate(imgs):
# 文件名处理。为每张图像生成一个默认文件名 image0 , image1 , ...。
f = f'image{i}' # filename
if isinstance(im, str): # filename or uri
# 如果图像是字符串(文件名或URI)。
# 如果 im 是字符串,检查它是否以 http 开头,如果是,则使用 requests 库下载图像并将其转换为NumPy数组;如果不是,则直接打开本地文件并转换。
# 更新文件名 f 为原始的文件名或URI。
im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im
elif isinstance(im, Image.Image): # PIL Image
# 处理PIL图像格式
# 如果 im 是PIL图像对象,将其转换为NumPy数组。 尝试获取PIL图像的文件名,如果没有则使用默认文件名。
im, f = np.asarray(im), getattr(im, 'filename', f) or f
# path.with_suffix(suffix)
# 在Python中, .with_suffix() 方法是 pathlib 模块中 Path 类的一个方法,它用于给文件路径添加或修改后缀。这个方法返回一个新的 Path 对象,其路径与原路径相同,但是文件的后缀被修改为指定的后缀。
# path : 一个 Path 对象,表示原始的文件路径。
# suffix : 一个新的后缀,可以包含前导点(如 .txt )或不包含(如 txt )。如果省略前导点, .with_suffix() 方法会自动添加。
# 返回值 :
# 这个方法返回一个新的 Path 对象,其路径与原路径相同,但是文件的后缀被修改为指定的后缀。
# 注意事项 :
# 如果原始路径没有后缀(即文件名中不包含点), .with_suffix() 方法将添加后缀,就像在文件名后直接添加一样。
# 如果 suffix 参数是空字符串, .with_suffix() 方法将移除文件的后缀,返回一个没有后缀的路径。
# 添加文件名到列表。将处理后的文件名添加到 files 列表中,确保文件名以 .jpg 为后缀。
files.append(Path(f).with_suffix('.jpg').name)
# 处理CHW格式的图像。
if im.shape[0] < 5: # image in CHW
# 如果图像的形状表明它是CHW格式(通道在前),则将其转换为HWC格式(高度、宽度、通道)。确保图像为三通道输入
im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
# 如果图像是三维的,确保它只有三个通道;如果是二维的(单通道),则复制通道以创建三通道图像。记录图像尺寸并调整
im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
# 记录图像尺寸并调整。记录原始图像的尺寸(高度和宽度)。
s = im.shape[:2] # HWC
shape0.append(s) # image shape
# 计算尺寸调整因子 g ,这是目标尺寸与图像最大尺寸的比值。 根据调整因子计算新的尺寸,并将其添加到 shape1 列表中
g = (size / max(s)) # gain
# 根据调整因子计算新的尺寸,并将其添加到 shape1 列表中。
shape1.append([y * g for y in s])
# 更新图像列表。将处理后的图像更新回原始图像列表 imgs 中。
imgs[i] = im # update
# 调整图像尺寸。
# def make_divisible(x, divisor): -> 它用于确保一个数值 x 能够被 divisor 整除。得到一个能够被 divisor 整除的数。 -> return math.ceil(x / divisor) * divisor
# np.stack(shape1, 0) 将 shape1 列表中的尺寸堆叠成一个 NumPy 数组。
# .max(0) 计算每个维度上的最大尺寸,得到一个包含最大宽度和高度的数组。
# make_divisible(x, int(self.stride.max())) 是一个自定义函数,它确保网络的输入尺寸可以被模型的步长整除,这通常是为了满足某些网络结构的要求。
# 列表推导式遍历最大尺寸,对每个尺寸应用 make_divisible 函数,得到最终的推理尺寸 shape1 。
shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
# def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
# -> 它用于对图像进行缩放和填充,以适应指定的尺寸,同时保持图像的宽高比。这种处理方式常用于深度学习中的目标检测和图像分类任务,以确保输入图像满足模型的尺寸要求。
# -> 函数返回 缩放和填充后的图像 img ,缩放比例 ratio ,以及宽度和高度的填充量 (dw, dh) 。
# -> return img, ratio, (dw, dh)
# letterbox 是一个自定义函数,它将图像填充到 shape1 指定的尺寸,保持图像的宽高比。 列表推导式遍历所有图像,对每张图像应用 letterbox 函数,并获取填充后的图像。
# auto=False 表示不自动调整图像尺寸,而是使用预先计算好的 shape1 。
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
# 堆叠图像。如果有多个图像( n > 1 ),则使用 np.stack 将它们堆叠成一个 NumPy 数组。 如果只有一张图像,则将其扩展维度以匹配批处理的格式。
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
# 转换图像格式。
# .transpose((0, 3, 1, 2)) 将图像从 BHWC(批处理、高度、宽度、通道)格式转换为 BCHW 格式。
# np.ascontiguousarray 确保数组在内存中是连续的,这对于某些深度学习框架的性能至关重要。转换数据类型并归一化
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
# 转换数据类型并归一化。
# torch.from_numpy 将 NumPy 数组转换为 PyTorch 张量。 .to(p.device) 将张量移动到模型所在的设备(CPU或GPU)。.type_as(p) 将张量的数据类型转换为模型参数的数据类型。 255. 归一化图像数据,将像素值从 [0, 255] 范围转换到 [0, 1] 范围。
x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
# def time_synchronized(): -> 它用于获取一个与CUDA时间同步的准确时间戳。 -> return time.time()
# 记录时间。调用 time_synchronized 函数记录当前时间,并将其添加到时间列表 t 中。
t.append(time_synchronized())
# 自动混合精度。
# amp.autocast 是PyTorch中的一个上下文管理器,用于自动优化模型的计算精度,以减少内存使用并提高计算速度。
# enabled=p.device.type != 'cpu' 指定当设备不是CPU时启用自动混合精度。
with amp.autocast(enabled=p.device.type != 'cpu'):
# Inference
# 模型推理。调用模型的前向传播方法 self.model ,传入预处理后的图像 x 和其他控制参数 augment 和 profile 。
# [0] 表示只获取模型输出的第一个元素,这通常是推理结果。
y = self.model(x, augment, profile)[0] # forward
# 记录推理时间。调用 time_synchronized 函数记录模型推理结束的时间,并将其添加到时间列表 t 中。
t.append(time_synchronized())
# Post-process 后处理。
# def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, labels=()):
# -> 该函数实现了非最大抑制(NMS)算法,用于在目标检测任务中去除重叠的检测框。返回包含NMS处理结果的输出张量 output 。
# -> return output
# 后处理 - 非极大值抑制(NMS)
# non_max_suppression 是一个函数,用于执行NMS,过滤掉重叠的边界框,只保留最佳的检测结果。 conf_thres 是置信度阈值, iou_thres 是交并比阈值, classes 是可选的类别过滤列表。
y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
# 坐标缩放。
for i in range(n):
# scale_coords 是一个函数,用于将NMS后的边界框坐标从推理尺寸缩放回原始图像尺寸。
# shape1 是推理尺寸, y[i][:, :4] 是第 i 张图像的边界框坐标, shape0[i] 是原始图像尺寸。
scale_coords(shape1, y[i][:, :4], shape0[i])
# 记录后处理时间。再次调用 time_synchronized 函数记录后处理结束的时间,并将其添加到时间列表 t 中。
t.append(time_synchronized())
# 返回检测结果。
# 创建一个 Detections 对象,包含原始图像列表 imgs 、推理结果 y 、文件名列表 files 、时间列表 t 、类别名称 self.names 和输入图像的形状 x.shape 。
# 返回这个对象,它包含了所有必要的信息,以便进一步处理或显示检测结果。
return Detections(imgs, y, files, t, self.names, x.shape)
# 这段代码定义了一个名为 Detections 的类,它用于存储和处理 YOLOv5 模型的推理结果。
# Detections 类用于封装 YOLOv5 模型的检测结果。
class Detections:
# YOLOv5 推理结果的检测类。
# detections class for YOLOv5 inference results
# 构造函数。
# 1.imgs :原始图像列表,通常为 NumPy 数组。
# 2.pred :模型预测结果,通常为张量列表,其中每个张量包含检测的边界框、置信度和类别。
# 3.files :图像文件名列表。
# 4.times :可选参数,用于存储推理过程中的时间戳。
# 5.names :可选参数,类别名称列表。
# 6.shape :可选参数,推理时的输入图像形状。
def __init__(self, imgs, pred, files, times=None, names=None, shape=None):
# 属性初始化。
# 调用父类的构造函数。
super(Detections, self).__init__()
# 获取预测结果张量所在的设备(CPU或GPU)。
d = pred[0].device # device
# 为每张图像创建一个归一化张量 gn ,包含图像的宽度、高度和两个1.0值(可能用于后续的归一化操作)。
gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
# 存储原始图像列表。
self.imgs = imgs # list of images as numpy arrays
# 存储模型预测结果。
self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
# 存储类别名称列表。
self.names = names # class names
# 存储图像文件名列表。
self.files = files # image filenames
# 存储 xyxy 格式的边界框坐标。
self.xyxy = pred # xyxy pixels
# def xyxy2xywh(x): -> 它用于将边界框的坐标从 xyxy 格式(即左上角和右下角的坐标)转换为 xywh 格式(即中心点坐标加上宽度和高度)。返回转换后的 xywh 格式的边界框坐标。 -> return y
# 将 xyxy 格式的边界框坐标转换为 xywh 格式。
self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
# 将 xyxy 格式的边界框坐标归一化。
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
# 将 xywh 格式的边界框坐标归一化。
self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
# 存储图像数量(批处理大小)。
self.n = len(self.pred) # number of images (batch size)
# 计算推理过程中的时间戳,并将时间转换为毫秒。
self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms)
# 存储推理时的输入图像形状。
self.s = shape # inference BCHW shape
# 这段代码定义了一个名为 display 的方法,它是 Detections 类的一个成员函数,用于处理和展示检测结果。
# 方法定义。
# 1.self :类的实例。
# 2.pprint :布尔值,如果为 True ,则打印检测结果的字符串描述。
# 3.show :布尔值,如果为 True ,则显示图像。
# 4.save :布尔值,如果为 True ,则保存图像。
# 5.render :布尔值,如果为 True ,则将渲染后的图像保存回 self.imgs 列表。
# 6.save_dir :字符串,指定保存图像的目录。
def display(self, pprint=False, show=False, save=False, render=False, save_dir=''):
# def color_list(): -> 它用于返回一个包含前10种 matplotlib 颜色表中的颜色的列表,每种颜色以 (r, g, b) 的形式表示。 -> return [hex2rgb(h) for h in matplotlib.colors.TABLEAU_COLORS.values()]
# 颜色列表。调用 color_list 函数获取颜色列表,用于绘制边界框。
colors = color_list()
# 遍历图像和预测结果。遍历图像列表 self.imgs 和预测结果列表 self.pred 。
for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
# 构建字符串描述。构建一个字符串,包含图像的索引和尺寸。
str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
if pred is not None:
# 如果预测结果 pred 不为空, pred[:, -1].unique() 获取所有唯一的类别索引。
for c in pred[:, -1].unique():
# 计算每个类别 c 的检测数量。
n = (pred[:, -1] == c).sum() # detections per class 每个类别的检测数。
# 将每个类别的 检测数量 和 类别名称 添加到字符串 str 中。如果检测数量大于1,则在类别名称后添加 's' 。
str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string 添加到字符串。
# 绘制边界框。
# 如果需要 显示 、 保存 或 渲染图像 ,则遍历预测结果 pred 。
if show or save or render:
# 解包每个预测结果,其中 box 是边界框坐标(x1, y1, x2, y2), conf 是置信度, cls 是类别索引。
for *box, conf, cls in pred: # xyxy, confidence, class
# 创建一个包含类别名称和置信度的标签。
label = f'{self.names[int(cls)]} {conf:.2f}'
# def plot_one_box(x, img, color=None, label=None, line_thickness=3):
# -> 它用于在图像上绘制单个边界框,并可选地添加标签。这个函数通常用于目标检测任务中,用于在检测到的对象周围绘制矩形框,并显示类别名称和置信度。
# 调用 plot_one_box 函数在图像 img 上绘制边界框和标签。 colors[int(cls) % 10] 从颜色列表 colors 中选择颜色,确保颜色索引在列表范围内。
plot_one_box(box, img, label=label, color=colors[int(cls) % 10])
# 转换图像格式。
# Image.fromarray(arr, mode=None)
# Image.fromarray 是 Python Imaging Library (PIL) 中的一个函数,用于从数组(通常是 NumPy 数组)创建图像。这个函数非常适合将数值数据转换为图像格式,这在图像处理和计算机视觉任务中非常常见。
# arr :输入的数组,通常是 NumPy 数组。数组的形状应该是 (height, width) 或 (height, width, channels) ,其中 channels 是颜色通道数。
# mode :可选参数,指定图像模式。如果不提供,PIL 将尝试从数组的数据类型推断模式。常见的模式包括 "L"(灰度图),"RGB"(红绿蓝彩色图),"RGBA"(红绿蓝透明度彩色图)等。
# 返回值 :这个函数返回一个 PIL Image 对象。
# 如果 img 是一个 NumPy 数组,那么使用 Image.fromarray 将其转换为 PIL 图像。这是因为 PIL 的 Image 模块可以方便地处理图像的显示和保存。
# img.astype(np.uint8) 确保图像数据类型为无符号8位整数,这是图像处理中常用的数据类型。
img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
# 打印检测信息。
if pprint:
# str.rstrip([chars])
# rstrip 是 Python 中字符串( str )对象的一个方法,用于移除字符串末尾的特定字符。如果没有指定字符,则默认移除空白字符(包括空格、制表符、换行符等)。
# chars :一个字符串,包含需要从末尾移除的字符集合。如果省略此参数,则默认移除空白字符。返回值这个方法返回一个新的字符串,原字符串末尾的指定字符被移除。
# 注意事项 :
# rstrip 方法不会修改原始字符串,而是返回一个新的字符串。
# 如果需要移除的字符在字符串中没有出现,则返回原始字符串的一个副本。
# 如果需要从字符串的开头移除字符,可以使用 lstrip 方法;如果需要从两端移除字符,可以使用 strip 方法。
# rstrip 是处理字符串时常用的方法之一,特别是在处理文件路径、用户输入或任何可能包含不需要的尾随字符的情况。
# 如果 pprint 参数为 True ,则打印变量 str 中的检测信息, str.rstrip(', ') 用于移除字符串末尾的逗号和空格。
print(str.rstrip(', '))
# 显示图像
if show:
# 如果 show 参数为 True ,则使用 PIL 图像的 show 方法显示图像。 self.files[i] 作为显示窗口的标题。
img.show(self.files[i]) # show
# 保存图像。
if save:
# 如果 save 参数为 True ,则使用 PIL 图像的 save 方法保存图像到指定的目录 save_dir 。
f = self.files[i]
# Path(save_dir) / f 构建完整的文件路径。
img.save(Path(save_dir) / f) # save
# 打印保存信息,如果 i 不是最后一个图像,则在文件名后添加逗号分隔符;如果是最后一个图像,则显示保存到的目录。
print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n')
# 渲染图像。
if render:
# 如果 render 参数为 True ,则将 PIL 图像转换回 NumPy 数组,并更新 self.imgs 列表中的相应元素。这允许后续处理或分析使用原始的 NumPy 数组格式。
self.imgs[i] = np.asarray(img)
# 这段代码定义了一个名为 print 的方法,它是 Detections 类的一个成员函数,用于打印检测结果和处理速度。
# 1.self :类的实例。
def print(self):
# 打印检测结果。调用 display 方法,并设置 pprint 参数为 True ,以便打印检测结果的字符串描述。
self.display(pprint=True) # print results
# 打印处理速度。
# 使用格式化字符串打印处理速度,包括预处理、推理和NMS(非极大值抑制)的时间,以及输入图像的形状。
# %.1fms 指定浮点数格式化为一位小数的毫秒数。
# self.t 是一个包含三个元素的元组,分别对应预处理、推理和NMS的时间(以毫秒为单位)。
# tuple(self.s) 将输入图像的形状 self.s 转换为元组,并将其作为字符串插入到打印的字符串中。
print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t)
# 这段代码定义了一个名为 show 的方法,它是 Detections 类的一个成员函数,用于显示检测结果。
def show(self):
# 调用 display 方法,并设置 show 参数为 True ,以便显示检测结果。这意味着 display 方法将处理图像,并使用适当的图形界面显示它们。
self.display(show=True) # show results
# 这段代码定义了一个名为 save 的方法,它是 Detections 类的一个成员函数,用于将检测结果保存到指定的目录。
# 1.self :类的实例。
# 2.save_dir :一个字符串参数,指定保存图像的目录,默认值为 'runs/hub/exp' 。
def save(self, save_dir='runs/hub/exp'):
# 增量保存目录。
# def increment_path(path, exist_ok=True, sep=''):
# -> 它用于生成一个唯一的文件路径。如果指定的路径已经存在,函数会通过添加一个数字后缀来创建一个新的路径。
# -> 如果路径存在且 exist_ok 为 True ,或者路径不存在,则直接返回原始路径。 / 如果路径存在且 exist_ok 为 False,返回一个新的路径,它是原始路径加上一个更新的数字后缀。
# -> return str(path) / return f"{path}{sep}{n}"
# increment_path 函数用于生成一个唯一的目录路径,如果指定的 save_dir 已经存在,则在目录名上添加一个数字后缀,以确保目录名唯一。
# exist_ok 参数指定是否允许覆盖现有目录,如果不是默认目录 'runs/hub/exp' ,则允许覆盖。
save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp') # increment save_dir
# 创建目录。使用 Path 函数创建一个目录对象。mkdir 方法创建目录, parents=True 允许创建多级目录, exist_ok=True 表示如果目录已存在则不抛出异常。
Path(save_dir).mkdir(parents=True, exist_ok=True)
# 保存检测结果。调用 display 方法,并设置 save 参数为 True ,以便保存检测结果。 save_dir 参数指定保存图像的目录。
self.display(save=True, save_dir=save_dir) # save results
# 这段代码定义了一个名为 render 的方法,它是 Detections 类的一个成员函数,用于渲染检测结果并将渲染后的图像保存回类的实例变量中。
# 1.self :类的实例。
def render(self):
# 渲染检测结果。调用 display 方法,并设置 render 参数为 True ,以便在图像上绘制边界框和标签,并将这些渲染后的图像保存回 self.imgs 列表。
self.display(render=True) # render results
# 返回渲染后的图像列表。返回更新后的 self.imgs 列表,其中包含了绘制了边界框和标签的图像。
return self.imgs
# 这段代码定义了一个名为 pandas 的方法,它是 Detections 类的一个成员函数,用于将检测结果转换为 pandas DataFrames,这样可以更方便地进行数据分析和处理。
# 1.self :类的实例。
def pandas(self):
# 将检测结果返回为 pandas DataFrames,即 print(results.pandas().xyxy[0])。
# return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0])
# 返回检测结果的副本。 使用 copy 函数创建 self 的副本,这样可以避免直接修改原始对象。
new = copy(self) # return copy
# 定义列名。 ca 和 cb 分别定义了 xyxy 和 xywh 格式的边界框数据的列名。
ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns
cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns
# 转换边界框数据。
# 使用 zip 函数遍历边界框数据的四种格式( xyxy , xyxyn , xywh , xywhn )和对应的列名。
# getattr(self, k) 获取 self 对象中与 k 对应的属性,即边界框数据。
# x.tolist() 将边界框数据转换为列表。
# x[:5] + [int(x[5]), self.names[int(x[5])]] 更新每个边界框数据,添加类别索引和类别名称。
# pd.DataFrame(x, columns=c) 将更新后的边界框数据转换为 pandas DataFrame。
# setattr(new, k, ...) 将转换后的 DataFrame 赋值给 new 对象中对应的属性。
for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]):
a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update
# setattr(object, name, value)
# setattr 是 Python 内置的一个函数,用于将属性赋值给对象。这个函数可以用来动态地设置对象的属性值,包括那些在代码运行时才知道名称的属性。
# object :要设置属性的对象。
# name :要设置的属性的名称,它应该是一个字符串。
# value :要赋给属性的值。
# 功能 :
# setattr 函数将 value 赋给 object 的 name 指定的属性。如果 name 指定的属性在 object 中不存在,则会创建一个新的属性。返回值 setattr 函数没有返回值。
# 注意事项 :
# 使用 setattr 时需要注意属性名称的字符串格式,因为属性名称会被直接用作对象的属性键。
# setattr 可以用于任何对象,包括自定义类的实例、内置类型的对象等。
# 如果需要删除对象的属性,可以使用 delattr 函数,其用法与 setattr 类似,但是用于删除属性而不是设置属性。
setattr(new, k, [pd.DataFrame(x, columns=c) for x in a])
# 返回新的 Detections 对象。 返回包含 pandas DataFrames 的 new 对象。
return new
# 这段代码定义了一个名为 tolist 的方法,它是 Detections 类的一个成员函数,用于将单个 Detections 对象中的检测结果拆分成多个单独的 Detections 对象,每个对象包含一张图像的检测结果。
# 1.self :类的实例。
def tolist(self):
# 返回 Detections 对象列表,即 'for result in results.tolist():'。
# return a list of Detections objects, i.e. 'for result in results.tolist():'
# 创建 Detections 对象列表
# 使用列表推导式创建一个新的 Detections 对象列表 x 。 对于 self.imgs 和 self.pred 中的每个元素(即每张图像及其对应的预测结果),创建一个新的 Detections 对象。
# self.imgs[i] 是第 i 张图像, [self.pred[i]] 是包含第 i 张图像预测结果的列表, self.names 是类别名称列表, self.s 是输入图像的形状。
x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)]
# 调整属性。
# 遍历新创建的 Detections 对象列表 x 。
for d in x:
# 对于每个对象 d ,遍历属性列表 ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn'] 。
for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
# getattr(object, name, default=None)
# getattr 是 Python 内置的一个函数,用于获取对象的属性值。这个函数可以用来动态地访问对象的属性,尤其是在属性名称在代码运行时才知道的情况下。
# object :要获取属性值的对象。
# name :要获取的属性的名称,它应该是一个字符串。
# default :可选参数,如果属性 name 在 object 中不存在时返回的默认值,默认为 None 。
# 功能 :
# getattr 函数返回 object 的 name 指定的属性值。如果 name 指定的属性在 object 中不存在,并且提供了 default 参数,则返回 default 指定的值;如果没有提供 default 参数,则抛出 AttributeError 异常。
# 返回值 :
# getattr 函数返回指定属性的值,或者在属性不存在时返回默认值。
# 注意事项 :
# 使用 getattr 时需要注意属性名称的字符串格式,因为属性名称会被直接用作对象的属性键。
# getattr 可以用于任何对象,包括自定义类的实例、内置类型的对象等。
# getattr 是一个非常有用的函数,特别是在处理动态属性名或需要在属性不存在时提供默认值的场景。
# 使用 getattr(d, k)[0] 获取属性 k 的值(这是一个列表),并取出列表中的第一个元素(因为每个 Detections 对象只包含一张图像的结果)。
# 使用 setattr(d, k, ...) 将取出的元素重新设置为属性 k 的值,从而将属性从列表中“弹出”。
setattr(d, k, getattr(d, k)[0]) # pop out of list
# 返回结果。
# 返回包含多个 Detections 对象的列表 x ,每个对象包含一张图像的检测结果。
return x
# 这段代码定义了一个名为 __len__ 的特殊方法(也称为魔术方法或内置方法),它在 Python 中用于让一个对象支持内置的 len() 函数。当使用 len() 函数时,Python 会自动调用这个方法。
def __len__(self):
# 这个方法返回实例变量 self.n 的值。通常, self.n 应该是一个整数,表示集合中的元素数量。
return self.n
# 使用场景 :
# 当你有一个自定义类,并且希望这个类的实例能够与 len() 函数一起使用时,你可以在类中定义 __len__ 方法。
# 注意事项 :
# __len__ 方法必须返回一个非负整数,因为长度不能是负数或非整数值。
# 如果你的类表示一个集合或序列,那么实现 __len__ 方法是很有用的,因为它允许用户使用 len() 函数来获取集合的大小。
# 如果 __len__ 方法不存在,尝试使用 len() 函数将导致 TypeError 。
# 这段代码定义了一个名为 Classify 的类,它是 nn.Module 的子类,用于构建一个分类头部(classification head),通常用于卷积神经网络(CNN)的末端,将特征图转换为分类结果。
# 类定义。 Classify 类继承自 PyTorch 的 nn.Module ,是一个神经网络模块。
# 构造函数。
# 1.c1 :输入通道数。
# 2.c2 :输出通道数,通常对应于分类类别的数量。
# 3.k :卷积核的大小,默认为 1。
# 4.s :卷积的步长,默认为 1。
# 5.p :填充大小,如果为 None ,则由 autopad 函数自动计算。
# 6.g :分组卷积的组数,默认为 1。
class Classify(nn.Module):
# 分类头,即 x(b,c1,20,20) 至 x(b,c2)。
# Classification head, i.e. x(b,c1,20,20) to x(b,c2)
def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
# 初始化。调用父类的构造函数。
super(Classify, self).__init__()
# 创建一个自适应平均池化层,将特征图的每个通道的空间维度(高和宽)减少到 1。
self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
# 创建一个卷积层,将输入通道数 c1 转换为输出通道数 c2 ,使用卷积核大小 k 和步长 s ,填充由 autopad 函数自动计算,组数为 g 。
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
# 创建一个展平层,用于将多维的张量展平成一维。
self.flat = nn.Flatten()
# 前向传播。
# 1.x :输入的特征图。
def forward(self, x):
# 如果输入 x 是一个列表,则对列表中的每个元素应用自适应平均池化,并将结果在维度 1 上拼接起来;如果 x 不是列表,则直接应用自适应平均池化。
z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
# 对拼接后的特征图 z 应用卷积,然后展平,得到最终的分类结果。
return self.flat(self.conv(z)) # flatten to x(b,c2)
# 例 : Classify 类的实例 classify_head = Classify(512, 1000) 将输入特征图 x = torch.randn(8, 512, 20, 20) 从形状 (8, 512, 20, 20) 转换为 (8, 1000) ,其中 8 是批量大小,1000 是类别数。
##### end of yolov5 ######
##### orepa #####
# transI_fusebn 函数的作用是将一个卷积核( kernel )和一个批量归一化层( bn )融合在一起。这种融合通常用于模型的推理阶段,以减少模型的计算复杂度和提高推理速度。
def transI_fusebn(kernel, bn):
# 获取批量归一化层的缩放因子( gamma ),它是一个可学习的参数,用于调整批量归一化后的输出。
gamma = bn.weight
# 计算批量归一化层的标准差( std )。这里, bn.running_var 是运行方差, bn.eps 是一个很小的数,用于防止除以零。 sqrt() 函数用于计算平方根。
std = (bn.running_var + bn.eps).sqrt()
# 这一行代码执行了两个操作。
# 首先,它将卷积核( kernel )与缩放因子除以标准差的结果相乘,并将结果重塑为一个四维张量( -1, 1, 1, 1 ),这样就可以与卷积核的形状匹配并进行元素乘法。
# 其次,它计算批量归一化层的偏移量( bn.bias )减去运行均值乘以缩放因子除以标准差的结果。这个操作是为了在融合后保持相同的输出分布。
# 最终,这个函数返回两个值 : 融合后的卷积核 和 调整后的偏移量 。
return kernel * ((gamma / std).reshape(-1, 1, 1, 1)), bn.bias - bn.running_mean * gamma / std
# 这样,原始的卷积和批量归一化操作就可以被一个单一的卷积操作所替代,从而在模型部署时减少计算量和内存占用。
# 这段代码定义了一个名为 ConvBN 的类,它是一个 PyTorch 神经网络模块,用于构建一个卷积层( Conv2d )后跟一个批量归一化层( BatchNorm2d )。这个类可以根据参数的不同配置来创建两种不同的结构:一种是部署模式( deploy ),另一种是非部署模式。
class ConvBN(nn.Module):
# 这是类的构造函数,用于初始化 ConvBN 实例。
# 1.in_channels :输入通道数。
# 2.out_channels :输出通道数。
# 3.kernel_size :卷积核的大小。
# 4.stride :卷积的步长,默认为1。
# 5.padding :卷积的填充,默认为0。
# 6.dilation :卷积的膨胀率,默认为1。
# 7.groups :分组卷积的组数,默认为1。
# 8.deploy :一个布尔值,指示是否处于部署模式。在部署模式下,不会使用批量归一化层。
# 9.nonlinear :非线性激活函数,默认为 nn.Identity() ,即没有激活函数。
def __init__(self, in_channels, out_channels, kernel_size,
stride=1, padding=0, dilation=1, groups=1, deploy=False, nonlinear=None):
# 这行代码调用了父类 nn.Module 的构造函数,是初始化 PyTorch 模块时的标准做法。
super().__init__()
# 这个条件判断检查 nonlinear 参数是否为 None 。
if nonlinear is None:
# 如果 nonlinear 为 None ,则将 self.nonlinear 设置为 nn.Identity() ,这是一个不改变输入的恒等激活函数,相当于没有激活函数。
self.nonlinear = nn.Identity()
else:
# 如果提供了 nonlinear 参数,则将 self.nonlinear 设置为传入的激活函数。
self.nonlinear = nonlinear
# 这个条件判断检查是否处于部署模式。
if deploy:
# 如果处于部署模式,创建一个带有偏置项( bias=True )的卷积层。这意味着不会使用批量归一化层,而是将卷积层的偏置项作为唯一的偏置。
self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
stride=stride, padding=padding, dilation=dilation, groups=groups, bias=True)
else:
# 如果不处于部署模式,创建一个不带偏置项( bias=False )的卷积层,并随后创建一个批量归一化层。卷积层的输出将通过这个批量归一化层进行归一化。
self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
stride=stride, padding=padding, dilation=dilation, groups=groups, bias=False)
self.bn = nn.BatchNorm2d(num_features=out_channels)
# 定义了 ConvBN 类的 forward 方法,它是 PyTorch 中神经网络模块的一个特殊方法,用于定义模块的前向传播逻辑。
def forward(self, x):
# 这个条件判断检查 self (即 ConvBN 实例)是否具有 bn 属性。这个属性是在构造函数中根据 deploy 参数的值来决定是否创建的。
if hasattr(self, 'bn'):
# 如果存在 bn 属性,说明不在部署模式下,因此会先通过卷积层 self.conv 处理输入 x ,然后将卷积的输出传递给批量归一化层 self.bn ,最后应用非线性激活函数 self.nonlinear 到结果上,并返回这个最终结果。
return self.nonlinear(self.bn(self.conv(x)))
else:
# 如果不存在 bn 属性,说明处于部署模式下,因此只会通过卷积层 self.conv 处理输入 x ,然后直接应用非线性激活函数 self.nonlinear 到卷积的输出上,并返回这个结果。
return self.nonlinear(self.conv(x))
# 在训练时,可以使用卷积层和批量归一化层的组合,而在部署时,可以省略批量归一化层,只使用卷积层,这样可以减少模型的复杂度和提高推理速度。
# ConvBN 类的 switch_to_deploy 方法,它用于将模块从训练模式转换到部署模式。在部署模式下,卷积层和批量归一化层将被融合为一个没有批量归一化层的卷积层。
# 定义了 switch_to_deploy 方法,这个方法不接受额外的参数,只作用于 self (即 ConvBN 实例)。
def switch_to_deploy(self):
# def transI_fusebn(kernel, bn):
# -> 将一个卷积核( kernel )和一个批量归一化层( bn )融合在一起。返回两个值 : 融合后的卷积核 和 调整后的偏移量 。
# -> return kernel * ((gamma / std).reshape(-1, 1, 1, 1)), bn.bias - bn.running_mean * gamma / std
# 调用 transI_fusebn 函数,将卷积层的权重 self.conv.weight 和批量归一化层 self.bn 融合,得到新的卷积核 kernel 和偏置 bias 。
kernel, bias = transI_fusebn(self.conv.weight, self.bn)
# 创建一个新的卷积层 conv ,其参数与原始卷积层 self.conv 相同,但 bias 参数设置为 True ,因为我们将使用融合后的偏置。
conv = nn.Conv2d(in_channels=self.conv.in_channels, out_channels=self.conv.out_channels, kernel_size=self.conv.kernel_size,
stride=self.conv.stride, padding=self.conv.padding, dilation=self.conv.dilation, groups=self.conv.groups, bias=True)
# 将融合后的卷积核和偏置分别赋值给新卷积层的权重和偏置。
conv.weight.data = kernel
conv.bias.data = bias
# 遍历 ConvBN 实例的所有参数,并调用 detach_ 方法。这会将参数从当前计算图中分离出来,使得在推理时不会跟踪梯度,这是部署模型时的常见做法。
for para in self.parameters():
para.detach_()
# 删除原始的 conv 和 bn 属性,这样在部署模式下,这些属性将不再存在。
self.__delattr__('conv')
self.__delattr__('bn')
# 将 ConvBN 实例的 conv 属性更新为新的卷积层,这样在部署模式下, forward 方法将只使用这个新的卷积层进行前向传播。
self.conv = conv
class OREPA_3x3_RepConv(nn.Module):
# 这是类的构造函数,用于初始化类的实例。
# 1.in_channels :输入通道数。
# 2.out_channels :输出通道数。
# 3.kernel_size :卷积核的大小。
# 4.stride=1 :卷积的步长,默认为1。
# 5.padding=0 :卷积的填充,默认为0。
# 6.dilation=1 :卷积的膨胀率,默认为1。
# 7.groups=1 :分组卷积的组数,默认为1。
# 8.internal_channels_1x1_3x3=None :这是一个可选参数,可能用于定义1x1和3x3卷积层之间的内部通道数。
# 9.deploy=False :一个布尔值,指示是否处于部署模式。
# 10.nonlinear=None :非线性激活函数,默认为 None 。
# 11.single_init=False :一个布尔值,可能用于指示是否使用单一初始化方法。
def __init__(self, in_channels, out_channels, kernel_size,
stride=1, padding=0, dilation=1, groups=1,
internal_channels_1x1_3x3=None,
deploy=False, nonlinear=None, single_init=False):
# 调用父类 nn.Module 的构造函数。
super(OREPA_3x3_RepConv, self).__init__()
# 将构造函数传入的 deploy 参数值赋给类的 deploy 属性。
self.deploy = deploy
if nonlinear is None:
# 如果传入的 nonlinear 参数为 None ,则将 self.nonlinear 设置为 nn.Identity() ,这是一个恒等激活函数,意味着它不对输入数据进行任何操作。
self.nonlinear = nn.Identity()
else:
# 如果提供了 nonlinear 参数,则将 self.nonlinear 设置为传入的激活函数。
self.nonlinear = nonlinear
# 将传入的 kernel_size 参数值赋给类的 kernel_size 属性,这个属性表示卷积核的大小。
self.kernel_size = kernel_size
# 将传入的 in_channels 参数值赋给类的 in_channels 属性,这个属性表示输入数据的通道数。
self.in_channels = in_channels
# 将传入的 out_channels 参数值赋给类的 out_channels 属性,这个属性表示输出数据的通道数。
self.out_channels = out_channels
# 将传入的 groups 参数值赋给类的 groups 属性,这个属性用于分组卷积,其中 groups=1 表示不使用分组卷积。
self.groups = groups
# 这是一个断言语句,用于确保 padding 参数的值等于 kernel_size 的一半。这种设置通常用于保持卷积操作的对称性,特别是在使用 stride=1 和 padding 等于核大小的一半时,可以保持特征图的空间维度不变。
assert padding == kernel_size // 2
# 断言语句 assert padding == kernel_size // 2 确保了卷积层的配置符合特定的设计要求,这在某些特定的网络架构中可能是必要的。
# 将传入的 stride 参数值赋给类的 stride 属性,这个属性表示卷积的步长。
self.stride = stride
# 将传入的 padding 参数值赋给类的 padding 属性,这个属性表示卷积的填充。
self.padding = padding
# 将传入的 dilation 参数值赋给类的 dilation 属性,这个属性表示卷积的膨胀率,用于控制卷积核中元素之间的间距。
self.dilation = dilation
# 初始化一个名为 branch_counter 的类属性,用于计数或追踪分支的数量。这个计数器的具体用途取决于类的其他部分,例如可能用于管理多个卷积分支或路径。
self.branch_counter = 0
# 创建一个新的 nn.Parameter 对象,它包含了一个四维的 torch.Tensor ,这个张量将被用作卷积层的权重。这个张量的维度是 (out_channels, in_channels//groups, kernel_size, kernel_size) 。
# 其中 : out_channels 是输出通道数。 int(in_channels/self.groups) 是每个组的输入通道数,这是分组卷积中每组处理的输入通道数。 kernel_size 是卷积核的尺寸,这里假设卷积核是正方形的。
self.weight_rbr_origin = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), kernel_size, kernel_size))
# torch.nn.init.kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu')
# nn.init.kaiming_uniform_ 是 PyTorch 中用于初始化神经网络权重的一种方法,它基于 Kaiming (He) 初始化,适用于使用 ReLU 或其变体(如 Leaky ReLU)作为激活函数的神经网络。
# 参数说明 :
# tensor :需要初始化的张量,通常是神经网络中的权重参数。
# a :负斜率系数,用于计算方差。当激活函数是 Leaky ReLU 时,该值为负斜率。默认值为 0,适用于 ReLU。
# mode :可以是 'fan_in' 或 'fan_out'。'fan_in' 模式保留前向传播的方差,'fan_out' 模式保留反向传播的方差。默认值为 'fan_in'。
# nonlinearity :非线性函数,建议仅与 'relu' 或 'leaky_relu'(默认)一起使用。
# 作用 :
# Kaiming 初始化的目的是在深层神经网络中保持每层输出的方差相对稳定,从而减少梯度消失或爆炸的问题,帮助模型更快地收敛。
# 使用 Kaiming 均匀初始化方法对权重参数 self.weight_rbr_origin 进行初始化。Kaiming 初始化(也称为 He 初始化)是一种专为ReLU激活函数设计的初始化方法,它可以减少训练深度神经网络时的初始梯度消失或爆炸问题。
# 参数 a 是均匀分布的缩放因子,这里设置为 math.sqrt(1.0) ,意味着权重将从一个均匀分布中采样,其范围是 [-sqrt(fan_in), sqrt(fan_in)] ,其中 fan_in 是权重张量的输入单元数。
nn.init.kaiming_uniform_(self.weight_rbr_origin, a=math.sqrt(1.0))
# 在创建了一个新的权重参数后,增加 branch_counter 的值。这可能表示类中增加了一个新的分支或路径,这个分支或路径将使用这个新初始化的权重参数。
self.branch_counter += 1
# 这个条件判断检查分组数 groups 是否小于输出通道数 out_channels 。如果是,那么执行以下操作。
if groups < out_channels:
# 创建两个 nn.Parameter 对象,它们包含了一维的 torch.Tensor ,这些张量将被用作额外的卷积层的权重。这些权重的维度是 (out_channels, in_channels//groups, 1, 1) ,表示它们是1x1卷积核。
self.weight_rbr_avg_conv = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), 1, 1))
self.weight_rbr_pfir_conv = nn.Parameter(torch.Tensor(out_channels, int(in_channels/self.groups), 1, 1))
# 使用 Kaiming 均匀初始化方法对这两个权重参数进行初始化。这里的 a 参数设置为 1.0 ,意味着权重将从一个均匀分布中采样,其范围是 [-sqrt(fan_in), sqrt(fan_in)] ,其中 fan_in 是权重张量的输入单元数。
nn.init.kaiming_uniform_(self.weight_rbr_avg_conv, a=1.0)
nn.init.kaiming_uniform_(self.weight_rbr_pfir_conv, a=1.0)
# 这两行代码似乎是多余的,因为它们只是访问了权重参数的数据,但没有进行任何操作。通常,这样的代码行不会出现在实际的类定义中,除非它们后面跟着对数据的操作。
self.weight_rbr_avg_conv.data
self.weight_rbr_pfir_conv.data
# torch.nn.Module.register_buffer(self, name, tensor)
# 在 PyTorch 中, register_buffer 是 nn.Module 类的一个方法,用于注册一个缓冲区(buffer)到模块中。缓冲区是模块状态的一部分,但它不会被优化器视为参数,也就是说,在模型训练过程中,缓冲区的内容不会被梯度更新。这通常用于存储不需要梯度的中间结果或者固定的参数。
# 参数说明 :
# self :当前的 nn.Module 实例。
# name :缓冲区的名称,这个名称必须是唯一的,并且不能与模块中的任何参数或子模块的名称冲突。
# tensor :要注册的缓冲区张量。
# 作用 :
# register_buffer 方法允许你将一个张量注册为模块的缓冲区,这样在保存和加载模型时,这个缓冲区也会被保存和加载。这对于需要在模型的不同部分之间共享数据,或者在模型的不同阶段使用固定数据的情况非常有用。
# 注册一个名为 weight_rbr_avg_avg 的缓冲区,它是一个二维张量,初始化为一个均匀的滤波器,其所有元素都是 1/kernel_size/kernel_size 。这个缓冲区可能用于平均池化或其他操作。
self.register_buffer('weight_rbr_avg_avg', torch.ones(kernel_size, kernel_size).mul(1.0/kernel_size/kernel_size))
# 在创建了额外的权重参数和缓冲区后,增加 branch_counter 的值。这可能表示类中增加了一个新的分支或路径,这个分支或路径将使用这些新初始化的权重参数和缓冲区。
self.branch_counter += 1
# 这段代码的目的是在 groups 参数小于 out_channels 的情况下,为 OREPA_3x3_RepConv 类添加额外的1x1卷积层的权重参数和一个用于平均操作的缓冲区。这些额外的参数可能用于实现特定的网络结构或功能,例如通道混合或特征融合。
# 如果 groups 大于或等于 out_channels ,则执行 else 代码块中的操作。
else:
# 在 else 代码块中,没有提供任何实现,而是直接抛出了一个 NotImplementedError 异常。这意味着当 groups 参数不小于 out_channels 时,当前的类或方法没有提供相应的实现,因此无法继续执行。
raise NotImplementedError
# 无论 groups < out_channels 条件是否满足, branch_counter 计数器都会增加 1。这可能表示构造函数的执行被视为创建了一个“分支”,或者是在类的初始化过程中增加了某种计数。
self.branch_counter += 1
# 这个条件判断检查 internal_channels_1x1_3x3 参数是否为 None 。如果是 None ,则意味着用户没有提供这个参数的值,需要在构造函数内部确定其值。
if internal_channels_1x1_3x3 is None:
# 这是一个条件表达式,用于根据 groups 和 out_channels 的关系来确定 internal_channels_1x1_3x3 的值。
# 如果 groups < out_channels ,则 internal_channels_1x1_3x3 被设置为 in_channels 。
# 否则, internal_channels_1x1_3x3 被设置为 2 * in_channels 。
# # For mobilenet, it is better to have 2X internal channels :这是一个注释,提供了为什么在 groups >= out_channels 时选择 2 * in_channels 的理由。在 MobileNet 架构中,通常更倾向于有两倍于输入通道数的内部通道数,这有助于保持网络的表达能力。
internal_channels_1x1_3x3 = in_channels if groups < out_channels else 2 * in_channels # For mobilenet, it is better to have 2X internal channels
# 这个条件判断检查 internal_channels_1x1_3x3 是否等于 in_channels 。
if internal_channels_1x1_3x3 == in_channels:
# 如果条件为真,创建一个全零的 nn.Parameter 对象,它包含了一个四维张量,这个张量将被用作1x1卷积层的权重。这个张量的维度是 (in_channels, in_channels//groups, 1, 1) 。
self.weight_rbr_1x1_kxk_idconv1 = nn.Parameter(torch.zeros(in_channels, int(in_channels/self.groups), 1, 1))
# 创建一个全零的 NumPy 数组,其形状与上面创建的权重张量相同。
id_value = np.zeros((in_channels, int(in_channels/self.groups), 1, 1))
# 遍历输入通道。
for i in range(in_channels):
# 在每个通道的对应位置上设置值为1,这样每个通道只与自己相连接,形成一个单位矩阵。
id_value[i, i % int(in_channels/self.groups), 0, 0] = 1
# 将 NumPy 数组转换为 PyTorch 张量,并确保其数据类型与 self.weight_rbr_1x1_kxk_idconv1 相同。
id_tensor = torch.from_numpy(id_value).type_as(self.weight_rbr_1x1_kxk_idconv1)
# 将 id_tensor 注册为类的缓冲区,这意味着它将被 PyTorch 跟踪,但不会被优化器优化。
self.register_buffer('id_tensor', id_tensor)
# 如果 internal_channels_1x1_3x3 不等于 in_channels ,则执行以下操作。
else:
# 创建一个新的 nn.Parameter 对象,它包含了一个四维张量,这个张量将被用作1x1卷积层的权重。这个张量的维度是 (internal_channels_1x1_3x3, in_channels//groups, 1, 1) 。
self.weight_rbr_1x1_kxk_conv1 = nn.Parameter(torch.Tensor(internal_channels_1x1_3x3, int(in_channels/self.groups), 1, 1))
# 使用 Kaiming 均匀初始化方法对 self.weight_rbr_1x1_kxk_conv1 进行初始化。
nn.init.kaiming_uniform_(self.weight_rbr_1x1_kxk_conv1, a=math.sqrt(1.0))
# 创建一个新的 nn.Parameter 对象,它包含了一个四维张量,这个张量将被用作3x3卷积层的权重。这个张量的维度是 (out_channels, internal_channels_1x1_3x3//groups, kernel_size, kernel_size) 。
self.weight_rbr_1x1_kxk_conv2 = nn.Parameter(torch.Tensor(out_channels, int(internal_channels_1x1_3x3/self.groups), kernel_size, kernel_size))
# 使用 Kaiming 均匀初始化方法对 self.weight_rbr_1x1_kxk_conv2 进行初始化。
nn.init.kaiming_uniform_(self.weight_rbr_1x1_kxk_conv2, a=math.sqrt(1.0))
# 在创建了新的权重参数后,增加 branch_counter 的值。
self.branch_counter += 1
# 定义了一个变量 expand_ratio ,其值为8。这个值用于扩展输入通道数,以便在深度可分离卷积中使用。
expand_ratio = 8
# 创建一个 nn.Parameter 对象,它包含了一个四维张量,这个张量将被用作深度卷积(Depthwise Convolution)的权重。这个张量的维度是 (in_channels*expand_ratio, 1, kernel_size, kernel_size) ,其中 in_channels*expand_ratio 是扩展后的输入通道数。
self.weight_rbr_gconv_dw = nn.Parameter(torch.Tensor(in_channels*expand_ratio, 1, kernel_size, kernel_size))
# 创建另一个 nn.Parameter 对象,它包含了一个四维张量,这个张量将被用作点卷积(Pointwise Convolution)的权重。这个张量的维度是 (out_channels, in_channels*expand_ratio, 1, 1) ,其中 in_channels*expand_ratio 是扩展后的输入通道数。
self.weight_rbr_gconv_pw = nn.Parameter(torch.Tensor(out_channels, in_channels*expand_ratio, 1, 1))
# 使用 Kaiming 均匀初始化方法对深度卷积的权重 self.weight_rbr_gconv_dw 进行初始化。
nn.init.kaiming_uniform_(self.weight_rbr_gconv_dw, a=math.sqrt(1.0))
# 使用 Kaiming 均匀初始化方法对点卷积的权重 self.weight_rbr_gconv_pw 进行初始化。
nn.init.kaiming_uniform_(self.weight_rbr_gconv_pw, a=math.sqrt(1.0))
# 在创建了新的权重参数后,增加 branch_counter 的值。这可能表示类中增加了一个新的分支或路径,这个分支或路径将使用这些新初始化的权重参数。
self.branch_counter += 1
# 这个条件判断检查输出通道数 out_channels 是否等于输入通道数 in_channels 并且步长 stride 是否等于1。
if out_channels == in_channels and stride == 1:
# 如果条件为真,即当输出通道数等于输入通道数且步长为1时,增加 branch_counter 的值。这可能表示在这种情况下,类的某个分支或路径被创建或激活。
self.branch_counter += 1
# 创建一个新的 nn.Parameter 对象,它包含了一个二维张量,这个张量将被用作类中的权重参数。这个张量的维度是 (self.branch_counter, self.out_channels) ,其中 self.branch_counter 是分支计数器的当前值, self.out_channels 是输出通道数。
self.vector = nn.Parameter(torch.Tensor(self.branch_counter, self.out_channels))
# 创建一个二维批量归一化层 nn.BatchNorm2d ,其参数是输出通道数 out_channels 。批量归一化是一种常用的技术,用于调整神经网络中间层的输出,以提高训练速度和性能。
self.bn = nn.BatchNorm2d(out_channels)
# 调用一个名为 fre_init 的自定义函数,这个函数可能是用于进一步初始化类的参数或执行其他一些初始化操作。
# def fre_init(self): -> 它是一个自定义的初始化函数,用于创建一个名为 prior_tensor 的张量,并将其注册为 OREPA_3x3_RepConv 类的一个缓冲区。
self.fre_init()
# torch.nn.init.constant_(tensor, val)
# nn.init.constant_ 是 PyTorch 中用于将张量(tensor)中的所有元素初始化为一个常数值的函数。这个函数会直接修改传入的张量,而不是返回一个新的张量。
# 参数说明 :
# tensor :需要被初始化的张量。
# val :一个常数值,张量中的所有元素都将被设置为这个值。
# 作用 :
# nn.init.constant_ 函数用于将张量中的每个元素设置为一个指定的常数值。这在某些情况下非常有用,比如当你需要将偏置项(bias)设置为特定的值,或者初始化某些参数为零(或其他值)以实现特定的网络行为时。
# 将 self.vector 参数的第一行(索引为0)的所有元素初始化为常数值0.25。这里的注释“origin”表明这一行可能与原始卷积路径的权重有关。
nn.init.constant_(self.vector[0, :], 0.25) #origin
# 将 self.vector 参数的第二行(索引为1)的所有元素初始化为常数值0.25。注释“avg”表明这一行可能与平均池化或某种平均操作的权重有关。
nn.init.constant_(self.vector[1, :], 0.25) #avg
# 将 self.vector 参数的第三行(索引为2)的所有元素初始化为常数值0.0。注释“prior”表明这一行可能与某种先验或额外的卷积路径的权重有关,这里初始化为0可能意味着这一路径在开始时不活跃。
nn.init.constant_(self.vector[2, :], 0.0) #prior
# 将 self.vector 参数的第四行(索引为3)的所有元素初始化为常数值0.5。注释“1x1_kxk”表明这一行可能与1x1和kxk卷积操作的权重有关。
nn.init.constant_(self.vector[3, :], 0.5) #1x1_kxk
# 将 self.vector 参数的第五行(索引为4)的所有元素初始化为常数值0.5。注释“dws_conv”表明这一行可能与深度可分离卷积操作的权重有关。
nn.init.constant_(self.vector[4, :], 0.5) #dws_conv
# self.vector 参数可能是一个权重向量,用于在不同的卷积路径或操作之间分配或调整权重。
# 通过使用 nn.init.constant_ 函数,这些路径或操作的初始权重被设置为特定的常数值,这可能是为了在训练开始时平衡不同路径的贡献,或者根据经验或理论分析预设某些路径的重要性。
# 这种初始化方法在设计复杂的卷积结构时很有用,特别是当需要结合多种不同的卷积操作或路径时,可以通过调整这些路径的权重来优化网络的性能。
# 这段代码定义了一个名为 fre_init 的方法,它是一个自定义的初始化函数,用于创建一个名为 prior_tensor 的张量,并将其注册为 OREPA_3x3_RepConv 类的一个缓冲区。
def fre_init(self):
# 创建一个新的 torch.Tensor ,其形状为 (out_channels, kernel_size, kernel_size) 。这个张量将被用于存储初始化的值。
prior_tensor = torch.Tensor(self.out_channels, self.kernel_size, self.kernel_size)
# 计算输出通道数的一半,并将其存储在变量 half_fg 中。这个值用于在后续的循环中区分两个不同的通道范围。
half_fg = self.out_channels/2
# 接下来的三个嵌套循环( for i in range(self.out_channels) , for h in range(3) , for w in range(3) )遍历 prior_tensor 的每个元素,并根据 i (通道索引)、 h (高度索引)和 w (宽度索引)的值计算每个元素的初始化值。
for i in range(self.out_channels):
for h in range(3):
for w in range(3):
# 如果当前的通道索引 i 小于 half_fg ,则使用公式计算 prior_tensor 中的值。
if i < half_fg:
# 公式基于余弦函数,其中 h+0.5 用于将卷积核的中心放在索引的中间。
prior_tensor[i, h, w] = math.cos(math.pi*(h+0.5)*(i+1)/3)
else:
# 如果当前的通道索引 i 大于或等于 half_fg ,则使用另一个公式计算 prior_tensor 中的值。
# 这个公式也是基于余弦函数,但是它使用 w+0.5 并将通道索引偏移了 half_fg 。
prior_tensor[i, h, w] = math.cos(math.pi*(w+0.5)*(i+1-half_fg)/3)
# 将 prior_tensor 注册为类的缓冲区,命名为 weight_rbr_prior 。这意味着 prior_tensor 将被保存在模型的状态中,并且在模型的参数字典中可见,但它不会被优化器更新。
self.register_buffer('weight_rbr_prior', prior_tensor)
# 这个方法的目的是使用余弦函数生成一个特定的初始化模式,这个模式可能与某些特定的网络结构或功能有关。
# 通过将这个张量注册为缓冲区,网络可以在训练过程中使用这个初始化模式,而不需要在每次迭代时重新计算它。这种初始化方法可能有助于网络学习特定的特征或模式。
# 定义了一个名为 weight_gen 的方法,它用于根据 OREPA_3x3_RepConv 类中的多个权重参数和向量 self.vector 生成最终的卷积权重。
# torch.einsum(eq, *operands)
# torch.einsum 是 PyTorch 提供的一个函数,它允许你使用爱因斯坦求和约定(Einstein summation convention)来指定复杂的张量运算。这个函数非常灵活,可以用于执行各种张量操作,包括转置、求和、点积、矩阵乘法等。
# 参数说明 :
# eq :一个字符串,指定了运算的方程式。这个字符串使用爱因斯坦求和约定来描述输入和输出张量的索引,以及如何对这些索引进行操作。
# operands :一个或多个张量,这些张量将被用于执行由 eq 指定的操作。
# 字符串格式 eq 字符串的格式如下 :
# 每个字符代表一个维度,例如, 'ijk' 表示三个维度。
# 重复的字符表示对该维度进行求和。
# 逗号( , )分隔不同的操作数。
# 箭头( -> )分隔操作数和结果。
# 这一行计算原始卷积权重 weight_rbr_origin 与 self.vector 的第一个元素(索引为0)的加权组合。这里使用的爱因斯坦求和约定(einsum notation) 'oihw,o->oihw' 表示。
# 'oihw' :权重张量的索引顺序,其中 o 代表输出通道, i 代表输入通道, h 和 w 分别代表高度和宽度。
# 'o' : self.vector 的索引,这里只考虑输出通道维度。
# ->oihw :输出张量的索引顺序,与权重张量相同。
weight_rbr_origin = torch.einsum('oihw,o->oihw', self.weight_rbr_origin, self.vector[0, :])
# 这一行首先计算平均卷积权重 weight_rbr_avg_conv 与平均滤波器 weight_rbr_avg_avg 的卷积,然后将结果与 self.vector 的第二个元素(索引为1)进行加权组合。
# 内部的 torch.einsum('oihw,hw->oihw', ...) 表示。
# 'oihw,hw->oihw' :将 weight_rbr_avg_conv 的空间维度(高度和宽度)与 weight_rbr_avg_avg 进行卷积。
weight_rbr_avg = torch.einsum('oihw,o->oihw', torch.einsum('oihw,hw->oihw', self.weight_rbr_avg_conv, self.weight_rbr_avg_avg), self.vector[1, :])
# 这一行首先计算先行卷积权重 weight_rbr_pfir_conv 与先行滤波器 weight_rbr_prior 的卷积,然后将结果与 self.vector 的第三个元素(索引为2)进行加权组合。
# 内部的 torch.einsum('oihw,ohw->oihw', ...) 表示。
# 'oihw,ohw->oihw' :将 weight_rbr_pfir_conv 的输出通道和空间维度(高度和宽度)与 weight_rbr_prior 进行卷积。
weight_rbr_pfir = torch.einsum('oihw,o->oihw', torch.einsum('oihw,ohw->oihw', self.weight_rbr_pfir_conv, self.weight_rbr_prior), self.vector[2, :])
# weight_rbr_1x1_kxk_conv1代表1x1卷积和kxk卷积组合的权重。
# 初始化 weight_rbr_1x1_kxk_conv1 为 None ,这是一个占位符,表示这个权重还没有被确定。
weight_rbr_1x1_kxk_conv1 = None
# 检查类实例中是否有名为 weight_rbr_1x1_kxk_idconv1 的属性。这个属性可能在类的某个分支中被创建,用于存储与单位矩阵相关的权重。
if hasattr(self, 'weight_rbr_1x1_kxk_idconv1'):
# torch.Tensor.squeeze([self, dim=None])
# .squeeze() 是 PyTorch 中的一个方法,用于去除张量(tensor)中所有长度为1的维度。这个方法返回一个新的张量,它与原始张量共享数据,但是没有长度为1的维度。
# 参数说明 :
# self :要进行操作的张量。
# dim=None :可选参数,指定要去除的维度。如果未指定或为 None ,则去除张量中所有长度为1的维度。
# 返回值 :
# 返回一个新的张量,其中所有长度为1的维度都被去除。
# 示例 :
# 原始张量 tensor 的形状是 (1, 3, 1, 5) ,使用 squeeze() 方法后,形状变为 (3, 5) ,因为长度为1的维度被去除了。
# 我们指定了 dim=0 ,所以只有第0维(长度为1)被去除,而其他长度为1的维度保持不变。
# .squeeze() 方法不会修改原始张量,而是返回一个新的张量。如果原始张量中没有长度为1的维度,或者指定的维度不是长度为1,那么返回的张量将与原始张量相同。
# 如果 weight_rbr_1x1_kxk_idconv1 存在,那么将这个属性与 id_tensor 相加,并将结果的任何单维度条目压缩( squeeze ),以去除长度为1的维度。 id_tensor 是一个之前注册的缓冲区,可能包含单位矩阵或其他与身份变换相关的权重。
weight_rbr_1x1_kxk_conv1 = (self.weight_rbr_1x1_kxk_idconv1 + self.id_tensor).squeeze()
# 如果 weight_rbr_1x1_kxk_conv1 不是 None ,检查类实例中是否有另一个名为 weight_rbr_1x1_kxk_conv1 的属性。这个属性可能在类的另一个分支中被创建,用于存储1x1卷积的权重。
elif hasattr(self, 'weight_rbr_1x1_kxk_conv1'):
# 如果这个属性存在,那么直接使用它,并调用 squeeze 方法来去除任何单维度条目。
weight_rbr_1x1_kxk_conv1 = self.weight_rbr_1x1_kxk_conv1.squeeze()
else:
# 如果上述两个属性都不存在,抛出 NotImplementedError 异常,表示这种情况没有被实现。
raise NotImplementedError
# 无论上述条件如何,都将 self.weight_rbr_1x1_kxk_conv2 赋值给 weight_rbr_1x1_kxk_conv2 。这个属性可能在类的构造函数中被初始化,并存储了kxk卷积的权重。
weight_rbr_1x1_kxk_conv2 = self.weight_rbr_1x1_kxk_conv2
# 当卷积层的分组数( groups )大于1时。 当 self.groups > 1 的情况。
if self.groups > 1:
# 将分组数赋值给变量 g 。
g = self.groups
# 获取 weight_rbr_1x1_kxk_conv1 张量的大小,并分别赋值给 t (张量的第0维大小)和 ig (张量的第1维大小)。
t, ig = weight_rbr_1x1_kxk_conv1.size()
# 获取 weight_rbr_1x1_kxk_conv2 张量的大小,并分别赋值给 o (张量的第0维大小)、 tg (张量的第1维大小)、 h (张量的第2维大小)和 w (张量的第3维大小)。
o, tg, h, w = weight_rbr_1x1_kxk_conv2.size()
# 将 weight_rbr_1x1_kxk_conv1 张量重新塑形为 (g, int(t/g), ig) 的形状,以适应分组卷积。
weight_rbr_1x1_kxk_conv1 = weight_rbr_1x1_kxk_conv1.view(g, int(t/g), ig)
# 将 weight_rbr_1x1_kxk_conv2 张量重新塑形为 (g, int(o/g), tg, h, w) 的形状,以适应分组卷积。
weight_rbr_1x1_kxk_conv2 = weight_rbr_1x1_kxk_conv2.view(g, int(o/g), tg, h, w)
# 使用 torch.einsum 函数计算分组卷积的组合权重,然后使用 .view(o, ig, h, w) 将结果重塑为 (o, ig, h, w) 的形状。
weight_rbr_1x1_kxk = torch.einsum('gti,gothw->goihw', weight_rbr_1x1_kxk_conv1, weight_rbr_1x1_kxk_conv2).view(o, ig, h, w)
else:
# 当分组数为1时,直接使用 torch.einsum 函数计算1x1和kxk卷积的组合权重。
weight_rbr_1x1_kxk = torch.einsum('ti,othw->oihw', weight_rbr_1x1_kxk_conv1, weight_rbr_1x1_kxk_conv2)
# 这一行使用 torch.einsum 函数将 weight_rbr_1x1_kxk 权重与 self.vector 的第四个元素(索引为3)进行加权组合。这里使用的爱因斯坦求和约定 'oihw,o->oihw' 表示对输出通道维度进行加权。
weight_rbr_1x1_kxk = torch.einsum('oihw,o->oihw', weight_rbr_1x1_kxk, self.vector[3, :])
# 这一行调用 dws2full 方法,将深度可分离卷积的权重 weight_rbr_gconv_dw 和 weight_rbr_gconv_pw 转换为完整的卷积权重。 self.in_channels 参数用于确保转换后的权重具有正确的输入通道数。
weight_rbr_gconv = self.dwsc2full(self.weight_rbr_gconv_dw, self.weight_rbr_gconv_pw, self.in_channels)
# 这一行再次使用 torch.einsum 函数将 weight_rbr_gconv 权重与 self.vector 的第五个元素(索引为4)进行加权组合。
weight_rbr_gconv = torch.einsum('oihw,o->oihw', weight_rbr_gconv, self.vector[4, :])
# 这一行将所有计算得到的权重相加,得到最终的卷积权重 weight 。这个总和包括原始 卷积权重 、 平均池化权重 、 1x1和kxk卷积权重 、 先行卷积权重 和 深度可分离卷积权重 。
# 最后一行返回 计算得到的最终卷积权重 。
weight = weight_rbr_origin + weight_rbr_avg + weight_rbr_1x1_kxk + weight_rbr_pfir + weight_rbr_gconv
return weight
# 这段代码定义了一个名为 dws2full 的方法,它将深度可分离卷积(Depthwise Separable Convolution)的权重转换为标准卷积(Full Convolution)的权重。
# 深度可分离卷积是一种优化技术,它将标准卷积分解为深度卷积(每个输入通道独立卷积)和点卷积(1x1卷积),以减少参数数量和计算量。
def dwsc2full(self, weight_dw, weight_pw, groups):
# 获取深度卷积权重 weight_dw 的大小,并分别赋值给 t (深度卷积的输出通道数,通常等于输入通道数乘以分组数)、 ig (每个分组的输入通道数)、 h (高度)和 w (宽度)。
t, ig, h, w = weight_dw.size()
# 获取点卷积权重 weight_pw 的大小,并分别赋值给 o (点卷积的输出通道数)和其他维度(这里不需要使用,所以使用 _ 占位)。
o, _, _, _ = weight_pw.size()
# 计算每个分组的深度卷积输出通道数。
tg = int(t/groups)
# 计算点卷积的输入通道数,通常是深度卷积输入通道数乘以分组数。
i = int(ig*groups)
# 将深度卷积权重 weight_dw 重新塑形为 (groups, tg, ig, h, w) 的形状,以适应分组和爱因斯坦求和约定。
weight_dw = weight_dw.view(groups, tg, ig, h, w)
# 将点卷积权重 weight_pw 去除单维度并重新塑形为 (o, groups, tg) 的形状。
weight_pw = weight_pw.squeeze().view(o, groups, tg)
# 使用 torch.einsum 函数将深度卷积权重和点卷积权重相乘,得到深度可分离卷积的权重。
# 这里使用的爱因斯坦求和约定 'gtihw,ogt->ogihw' 表示 :
# 'gtihw' :深度卷积权重的索引顺序。
# 'ogt' :点卷积权重的索引顺序。
# 'ogihw' :输出权重的索引顺序。
weight_dsc = torch.einsum('gtihw,ogt->ogihw', weight_dw, weight_pw)
# 将深度可分离卷积权重 weight_dsc 重新塑形为 (o, i, h, w) 的形状,并返回。这个形状与标准卷积权重的形状相同。
return weight_dsc.view(o, i, h, w)
# 这个方法的目的是将深度可分离卷积的权重转换为标准卷积权重,以便在需要时可以在网络中使用标准卷积。这种转换在某些网络结构中很有用,例如在需要将深度可分离卷积与其他类型的卷积结合使用时。
# 这段代码定义了 OREPA_3x3_RepConv 类的 forward 方法,它指定了模型在前向传播时的行为。
def forward(self, inputs):
# 调用 weight_gen 方法来生成当前层的卷积权重。 weight_gen 方法包含复杂的逻辑,用于结合不同的权重参数和可能的动态权重调整。
weight = self.weight_gen()
# 使用 PyTorch 的 F.conv2d 函数执行卷积操作。这里, inputs 是输入数据, weight 是在上一步生成的权重, bias 设置为 None 表示不使用偏置项, stride 、 padding 、 dilation 和 groups 分别是卷积的步长、填充、膨胀率和分组数,这些参数都是从类的属性中获取的。
out = F.conv2d(inputs, weight, bias=None, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=self.groups)
# 首先,对卷积输出 out 应用批量归一化 self.bn ,然后应用非线性激活函数 self.nonlinear 。 self.nonlinear 可能是一个如 nn.ReLU 、 nn.Sigmoid 或 nn.Identity (恒等激活,即不改变输入)的激活函数。
# 最终,返回 经过激活函数处理后的输出 。
return self.nonlinear(self.bn(out))
# 定义了 RepConv_OREPA 类,它继承自 nn.Module 。
class RepConv_OREPA(nn.Module):
# 1. c1 :输入通道数( in_channels )。
# 2. c2 :输出通道数( out_channels )。
# 3. k=3 :卷积核的大小,默认为3。
# 4. s=1 :卷积的步长(stride),默认为1。
# 5. padding=1 :卷积的填充(padding),默认为1。这通常用于保持特征图的尺寸。
# 6. dilation=1 :卷积的膨胀率(dilation),默认为1。膨胀卷积允许扩大卷积核的接收区域。
# 7. groups=1 :分组卷积的组数(groups)。当 groups 等于输入通道数时,每个输入通道都有自己的卷积分支。
# 8. padding_mode='zeros' :填充模式,默认为 'zeros',表示使用零填充。
# 9. deploy=False :一个布尔值,指示是否处于部署模式。在部署模式下,可能会有一些优化,例如融合卷积和批量归一化层。
# 10. use_se=False :一个布尔值,指示是否使用SE(Squeeze-and-Excitation)模块,这是一种注意力机制,用于增强特征的表示能力。
# 11. nonlinear=nn.SiLU() :使用的非线性激活函数,默认为SiLU(也称为Swish),这是一种自门控的激活函数。
def __init__(self, c1, c2, k=3, s=1, padding=1, dilation=1, groups=1, padding_mode='zeros', deploy=False, use_se=False, nonlinear=nn.SiLU()):
super(RepConv_OREPA, self).__init__()
# 将构造函数传入的 deploy 参数值赋给类的 deploy 属性,用于指示是否处于部署模式。
self.deploy = deploy
# 将构造函数传入的 groups 参数值赋给类的 groups 属性,这个属性用于分组卷积。
self.groups = groups
# 将构造函数传入的 c1 参数值赋给类的 in_channels 属性,这个属性表示输入数据的通道数。
self.in_channels = c1
# 将构造函数传入的 c2 参数值赋给类的 out_channels 属性,这个属性表示输出数据的通道数。
self.out_channels = c2
# 将构造函数传入的 padding 参数值赋给类的 padding 属性,这个属性表示卷积的填充。
self.padding = padding
# 将构造函数传入的 dilation 参数值赋给类的 dilation 属性,这个属性表示卷积的膨胀率。
self.dilation = dilation
self.groups = groups
# 断言语句,确保卷积核大小 k 等于3。如果 k 不等于3,将抛出一个 AssertionError 。
assert k == 3
# 断言语句,确保卷积的填充 padding 等于1。如果 padding 不等于1,将抛出一个 AssertionError 。
assert padding == 1
# 计算用于1x1卷积的填充值 padding_11 。由于卷积核大小 k 为3,这个计算将确保1x1卷积的填充是对称的,即 padding_11 = 1 - 3 // 2 = 1 - 1 = 0 。
padding_11 = padding - k // 2
# 检查是否提供了非线性激活函数 nonlinear 。
if nonlinear is None:
# 如果 nonlinear 为 None ,则将 self.nonlinearity 设置为 nn.Identity() ,这是一个恒等激活函数,意味着它不对输入数据进行任何操作。
self.nonlinearity = nn.Identity()
else:
# 如果提供了 nonlinear 参数,则将 self.nonlinearity 设置为传入的激活函数。
self.nonlinearity = nonlinear
# 检查是否要使用SE(Squeeze-and-Excitation)模块。
if use_se:
# 如果 use_se 为 True ,则创建一个 SEBlock 实例,并将其赋值给 self.se 。 SEBlock 是一个自定义的类,用于实现SE模块的功能。 internal_neurons 参数设置为输出通道数的1/16,这是SE模块中常见的设计。
self.se = SEBlock(self.out_channels, internal_neurons=self.out_channels // 16)
else:
# 如果 use_se 为 False ,则将 self.se 设置为 nn.Identity() ,意味着不使用SE模块,而是使用恒等变换。
self.se = nn.Identity()
# 当 deploy 为 True 时。
if deploy:
# 创建一个标准的 nn.Conv2d 卷积层,用于部署模式。这个卷积层将所有的操作(包括批量归一化和激活函数)融合到一个卷积操作中,以便于模型的高效推理。
# 这里的参数包括输入通道数 self.in_channels 、输出通道数 self.out_channels 、卷积核大小 k 、步长 s 、填充 padding 、膨胀率 dilation 、分组数 groups ,以及 bias 和 padding_mode 。
self.rbr_reparam = nn.Conv2d(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=k, stride=s,
padding=padding, dilation=dilation, groups=groups, bias=True, padding_mode=padding_mode)
# 当 deploy 为 False 时。
else:
# 如果输出通道数等于输入通道数且步长为1,创建一个 nn.BatchNorm2d 批量归一化层作为 self.rbr_identity 。这可能用于实现残差连接中的恒等映射。如果不满足条件,则 self.rbr_identity 被设置为 None 。
self.rbr_identity = nn.BatchNorm2d(num_features=self.in_channels) if self.out_channels == self.in_channels and s == 1 else None
# 创建一个 OREPA_3x3_RepConv 实例,用于实现3x3的重复卷积结构。这个类可能包含了更多的复杂逻辑,用于实现重复卷积和可能的注意力机制。
self.rbr_dense = OREPA_3x3_RepConv(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=k, stride=s, padding=padding, groups=groups, dilation=1)
# 创建一个 ConvBN 实例,用于实现1x1的卷积结构,这个结构可能包括卷积后接批量归一化。
self.rbr_1x1 = ConvBN(in_channels=self.in_channels, out_channels=self.out_channels, kernel_size=1, stride=s, padding=padding_11, groups=groups, dilation=1)
# 打印一条消息,显示当前 RepVGG 块的 identity 连接的状态。
print('RepVGG Block, identity = ', self.rbr_identity)
# 这段代码定义了 RepConv_OREPA 类的 forward 方法,它指定了模型在前向传播时的行为。
def forward(self, inputs):
# 检查是否存在名为 rbr_reparam 的属性。这个属性仅在部署模式下创建,它是一个融合了所有操作(卷积、批量归一化、激活函数等)的卷积层。
if hasattr(self, 'rbr_reparam'):
# 如果存在 rbr_reparam ,则使用这个融合层处理输入 inputs ,然后应用 SE(Squeeze-and-Excitation)模块和非线性激活函数 self.nonlinearity ,最后返回处理后的结果。
return self.nonlinearity(self.se(self.rbr_reparam(inputs)))
# 检查是否没有恒等映射( rbr_identity )。恒等映射仅在输出通道数等于输入通道数且步长为1时创建。
if self.rbr_identity is None:
# 如果没有恒等映射,则 id_out 设置为0。
id_out = 0
else:
# 如果有恒等映射,则计算恒等映射的输出。
id_out = self.rbr_identity(inputs)
# 计算 rbr_dense (3x3重复卷积)的输出。
out1 = self.rbr_dense(inputs)
# 计算 rbr_1x1 (1x1卷积)的输出。
out2 = self.rbr_1x1(inputs)
# 将恒等映射的输出赋值给 out3 。
out3 = id_out
# 将三个输出相加,得到最终的输出。这里的加法操作实现了残差连接,其中 out1 和 out2 是卷积分支的输出, out3 是恒等映射的输出。
out = out1 + out2 + out3
# 对最终的输出应用 SE 模块和非线性激活函数 self.nonlinearity ,然后返回处理后的结果。
return self.nonlinearity(self.se(out))
# 这个方法展示了 RepConv_OREPA 类如何灵活地处理不同的网络配置,同时支持残差连接和注意力机制。在非部署模式下,它通过多个分支处理输入,然后将这些分支的输出相加,最后应用 SE 模块和非线性激活函数。这种设计允许模型在训练和部署时具有不同的行为,以适应不同的需求。
# Optional. This improves the accuracy and facilitates quantization.
# 1. Cancel the original weight decay on rbr_dense.conv.weight and rbr_1x1.conv.weight.
# 2. Use like this.
# loss = criterion(....)
# for every RepVGGBlock blk:
# loss += weight_decay_coefficient * 0.5 * blk.get_cust_L2()
# optimizer.zero_grad()
# loss.backward()
# Not used for OREPA 不用于 OREPA。
# 这段代码定义了 RepConv_OREPA 类的 get_custom_L2 方法,它用于计算自定义的 L2 损失,这种损失可能用于正则化网络的权重。
def get_custom_L2(self):
# 调用 self.rbr_dense (3x3重复卷积模块)的 weight_gen 方法来获取其权重。
K3 = self.rbr_dense.weight_gen()
# 获取 self.rbr_1x1 (1x1卷积模块)的权重。
K1 = self.rbr_1x1.conv.weight
# 计算 self.rbr_dense 的批量归一化层的缩放因子,并将其重塑为四维张量,然后从计算图中分离( detach ),使其不会影响梯度。
t3 = (self.rbr_dense.bn.weight / ((self.rbr_dense.bn.running_var + self.rbr_dense.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()
# 同上,但是针对 self.rbr_1x1 的批量归一化层。
t1 = (self.rbr_1x1.bn.weight / ((self.rbr_1x1.bn.running_var + self.rbr_1x1.bn.eps).sqrt())).reshape(-1, 1, 1, 1).detach()
# 计算3x3卷积核的L2损失,但是排除中心点(索引为 [1:2, 1:2] 的元素),因为中心点将被用于计算等效的中心点。
l2_loss_circle = (K3 ** 2).sum() - (K3[:, :, 1:2, 1:2] ** 2).sum() # The L2 loss of the "circle" of weights in 3x3 kernel. Use regular L2 on them. 3x3 内核中权重“圆”的 L2 损失。对其使用常规 L2。
# 计算3x3卷积核的中心点与1x1卷积核的加权和,这个加权和代表了3x3卷积核中心点的等效结果。
eq_kernel = K3[:, :, 1:2, 1:2] * t3 + K1 * t1 # The equivalent resultant central point of 3x3 kernel. 等效于 3x3 核的合成中心点。
# 计算等效中心点的L2损失,并对其进行归一化,使得L2系数与常规L2损失可比。
l2_loss_eq_kernel = (eq_kernel ** 2 / (t3 ** 2 + t1 ** 2)).sum() # Normalize for an L2 coefficient comparable to regular L2. 对与常规 L2 相当的 L2 系数进行标准化。
# 返回 等效中心点的L2损失 和 3x3卷积核的“圈”(即非中心点)的L2损失 之和。
return l2_loss_eq_kernel + l2_loss_circle
# 这段代码定义了 RepConv_OREPA 类的 get_equivalent_kernel_bias 方法,它用于获取等效的卷积核和偏置,这些等效的参数考虑了网络中的不同分支(如3x3卷积、1x1卷积和恒等映射)以及批量归一化层的影响。
def get_equivalent_kernel_bias(self):
# def _fuse_bn_tensor(self, branch):
# -> 它用于将批量归一化层( BatchNorm2d )的参数融合到卷积核中。这种融合通常用于模型部署,以减少模型在推理时的计算量。返回 融合后的卷积核 和 偏置 。
# ->return kernel * t, beta - running_mean * gamma / std
# 调用 _fuse_bn_tensor 方法处理 rbr_dense 分支(3x3卷积),获取等效的卷积核 kernel3x3 和偏置 bias3x3 。
kernel3x3, bias3x3 = self._fuse_bn_tensor(self.rbr_dense)
# 调用 _fuse_bn_tensor 方法处理 rbr_1x1 分支(1x1卷积),获取等效的卷积核 kernel1x1 和偏置 bias1x1 。
kernel1x1, bias1x1 = self._fuse_bn_tensor(self.rbr_1x1)
# 调用 _fuse_bn_tensor 方法处理 rbr_identity 分支(恒等映射),获取等效的卷积核 kernelid 和偏置 biasid 。
kernelid, biasid = self._fuse_bn_tensor(self.rbr_identity)
# 将三个分支的等效卷积核和偏置相加,得到 最终的等效卷积核 和 偏置 。这里, self._pad_1x1_to_3x3_tensor(kernel1x1) 方法用于将1x1卷积核填充为3x3卷积核,以便可以与其它3x3卷积核相加。
return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
# 这个方法的目的是提供一个方式来获取网络中所有相关分支的等效卷积核和偏置,这对于分析网络的行为、进行网络剪枝或者在部署时融合网络层特别有用。通过这种方式,可以将复杂的多分支网络结构简化为一个等效的单分支结构,同时保留原始网络的行为。
# 这段代码定义了一个名为 _pad_1x1_to_3x3_tensor 的方法,它用于将1x1卷积核通过填充(padding)转换为3x3卷积核。
# 1.kernel1x1 :代表1x1卷积核的权重。
def _pad_1x1_to_3x3_tensor(self, kernel1x1):
# 检查传入的1x1卷积核是否为 None 。
if kernel1x1 is None:
# 如果 kernel1x1 是 None ,则返回0。这可能表示没有1x1卷积核需要处理,或者在某些条件下1x1卷积核不被使用。
return 0
# 如果 kernel1x1 不是 None ,则执行以下操作。
else:
# torch.nn.functional.pad(input, pad, mode='xxx', value=x)
# torch.nn.functional.pad 是 PyTorch 中用于对张量进行填充操作的函数。填充操作在处理图像、序列数据等任务时非常常见,它可以在张量的指定维度两端添加一定数量的元素,填充方式多样,包括零填充、常数填充、反射填充和边界填充等。
# 返回一个新的张量,对输入张量进行了指定方式的填充。
# input (Tensor) :输入的张量。
# pad (tuple) :指定每个维度填充的数目。格式为 (left, right, top, bottom, …)。
# mode (str, 可选) :填充模式,包括 constant (常数填充,默认)、 reflect (反射填充)、 replicate (边界填充)和 circular( 循环填充)。
# value (float, 可选) :常数填充模式下填充值,仅当 mode 为 ‘ constant ’ 时有效,其他模式时即使指定了值也会被忽略掉。
# 使用 PyTorch 的 torch.nn.functional.pad 函数对1x1卷积核进行填充。填充的参数 [1,1,1,1] 表示在卷积核的上下左右各填充1个单位的零。这样,一个1x1的卷积核就被扩展为3x3的卷积核,其中新增的元素都是0。
return torch.nn.functional.pad(kernel1x1, [1,1,1,1])
# 这个方法的目的是将1x1卷积核转换为3x3卷积核,这在某些卷积网络的设计中可能是必要的,例如在实现深度可分离卷积(Depthwise Separable Convolution)时,需要将1x1卷积核与3x3深度卷积核结合使用。
# 通过这种方法,可以在不改变卷积核权重分布的情况下,将1x1卷积核适配到需要3x3卷积核的场景中。
# 这段代码定义了一个名为 _fuse_bn_tensor 的方法,它用于将批量归一化层( BatchNorm2d )的参数融合到卷积核中。这种融合通常用于模型部署,以减少模型在推理时的计算量。
# branch :它代表一个分支,可能是一个批量归一化层或者包含卷积和批量归一化的复合层。
def _fuse_bn_tensor(self, branch):
# 检查传入的分支是否为 None 。
if branch is None:
# 如果 branch 是 None ,则返回两个0值。这可能表示没有分支需要处理,或者在某些条件下分支不被使用。
return 0, 0
# 检查 branch 是否不是 nn.BatchNorm2d 类型。
if not isinstance(branch, nn.BatchNorm2d):
# 检查 branch 是否是 OREPA_3x3_RepConv 类型。
if isinstance(branch, OREPA_3x3_RepConv):
# 如果是 OREPA_3x3_RepConv 类型,调用其 weight_gen 方法来获取卷积核。
kernel = branch.weight_gen()
# 检查 branch 是否是 ConvBN 类型。
elif isinstance(branch, ConvBN):
# 如果是 ConvBN 类型,直接获取其卷积核。
kernel = branch.conv.weight
# 如果 branch 不是上述任何一种类型。
else:
# 抛出 NotImplementedError 异常,表示这种情况没有被实现。
raise NotImplementedError
# 获取 branch 的批量归一化层的运行均值。
running_mean = branch.bn.running_mean
# 获取 branch 的批量归一化层的运行方差。
running_var = branch.bn.running_var
# 获取 branch 的批量归一化层的缩放因子(也称为权重)。
gamma = branch.bn.weight
# 获取 branch 的批量归一化层的偏移量(也称为偏置)。
beta = branch.bn.bias
# 获取 branch 的批量归一化层的epsilon值,这是一个小的常数,用于防止除以零。
eps = branch.bn.eps
# branch 是 nn.BatchNorm2d 类型。
else:
# 检查类实例中是否没有名为 id_tensor 的属性。
if not hasattr(self, 'id_tensor'):
# 计算每个分组的输入通道数。
input_dim = self.in_channels // self.groups
# 创建一个全零的 NumPy 数组,其形状为 (in_channels, input_dim, 3, 3) ,用于存储单位矩阵形式的卷积核。
kernel_value = np.zeros((self.in_channels, input_dim, 3, 3), dtype=np.float32)
# 遍历输入通道。
for i in range(self.in_channels):
# 在每个通道的对应位置上设置值为1,这样每个通道只与自己相连接,形成一个单位矩阵。
kernel_value[i, i % input_dim, 1, 1] = 1
# 将 NumPy 数组转换为 PyTorch 张量,并确保它被移动到与 branch.weight 相同的设备上,然后注册为 id_tensor 属性。
self.id_tensor = torch.from_numpy(kernel_value).to(branch.weight.device)
# 将 id_tensor 赋值给 kernel 变量。
kernel = self.id_tensor
# 获取 branch 的批量归一化层的运行均值。
running_mean = branch.running_mean
# 获取 branch 的批量归一化层的运行方差。
running_var = branch.running_var
# 获取 branch 的批量归一化层的缩放因子(也称为权重)。
gamma = branch.weight
# 获取 branch 的批量归一化层的偏移量(也称为偏置)。
beta = branch.bias
# 获取 branch 的批量归一化层的epsilon值。
eps = branch.eps
# 计算批量归一化层的标准差。
std = (running_var + eps).sqrt()
# 计算调整后的缩放因子,然后重塑为四维张量。
t = (gamma / std).reshape(-1, 1, 1, 1)
# 返回 融合后的卷积核 和 偏置 。卷积核是通过将单位矩阵卷积核与调整后的缩放因子相乘得到的,偏置是通过从原始偏置中减去均值与缩放因子的乘积得到的。
return kernel * t, beta - running_mean * gamma / std
# 这个方法的目的是将批量归一化层的参数融合到卷积核中,以便在模型部署时可以省略批量归一化层,直接使用融合后的卷积核进行计算。这种优化可以减少模型的计算量和内存占用,提高推理速度。通过这种方式,可以在不牺牲模型性能的情况下,提高模型的运行效率。
# 这段代码定义了 RepConv_OREPA 类的 switch_to_deploy 方法,它用于将模型从训练模式转换到部署模式。在部署模式下,模型会将多个卷积和批量归一化层融合为一个单一的卷积层,以提高推理效率。
def switch_to_deploy(self):
# 检查是否已经存在一个名为 rbr_reparam 的属性,这意味着模型已经在部署模式下。
if hasattr(self, 'rbr_reparam'):
# 如果已经在部署模式下,则直接返回,不执行任何操作。
return
# 打印一条消息,表示正在将 RepConv_OREPA 转换到部署模式。
print(f"RepConv_OREPA.switch_to_deploy")
# def get_equivalent_kernel_bias(self):
# -> 它用于获取等效的卷积核和偏置,这些等效的参数考虑了网络中的不同分支(如3x3卷积、1x1卷积和恒等映射)以及批量归一化层的影响。将三个分支的等效卷积核和偏置相加,得到 最终的等效卷积核 和 偏置 。
# -> return kernel3x3 + self._pad_1x1_to_3x3_tensor(kernel1x1) + kernelid, bias3x3 + bias1x1 + biasid
# 调用 get_equivalent_kernel_bias 方法来获取等效的卷积核和偏置,这些参数考虑了所有相关的分支(如3x3卷积、1x1卷积和恒等映射)。
kernel, bias = self.get_equivalent_kernel_bias()
# 创建一个新的 nn.Conv2d 卷积层,并将等效的卷积核和偏置赋值给它。这个新的卷积层包含了原始模型中所有卷积和批量归一化层的信息。
self.rbr_reparam = nn.Conv2d(in_channels=self.rbr_dense.in_channels, out_channels=self.rbr_dense.out_channels,
kernel_size=self.rbr_dense.kernel_size, stride=self.rbr_dense.stride,
padding=self.rbr_dense.padding, dilation=self.rbr_dense.dilation, groups=self.rbr_dense.groups, bias=True)
# 将等效的卷积核赋值给 rbr_reparam 的权重。
self.rbr_reparam.weight.data = kernel
# 将等效的偏置赋值给 rbr_reparam 的偏置。
self.rbr_reparam.bias.data = bias
# 遍历模型的所有参数。
for para in self.parameters():
# 将每个参数从计算图中分离,这样它们就不会在后续的计算中跟踪梯度。
para.detach_()
# 删除 rbr_dense 属性,它包含了原始的3x3卷积层。
self.__delattr__('rbr_dense')
# 删除 rbr_1x1 属性,它包含了原始的1x1卷积层。
self.__delattr__('rbr_1x1')
# 检查是否存在 rbr_identity 属性。
if hasattr(self, 'rbr_identity'):
# 如果存在,删除 rbr_identity 属性,它包含了原始的恒等映射层。
self.__delattr__('rbr_identity')
##### end of orepa #####
##### swin transformer #####
# Swin Transformer 是一种为计算机视觉任务设计的分层Transformer结构,由微软亚洲研究院提出,并在2021年的ICCV会议上获得了最佳论文奖。
# 它能够有效地作为多种视觉任务的通用骨干网络,如图像分类、目标检测和语义分割等。
# Swin Transformer的主要特点包括 :
# 1. 分层结构 (Hierarchical Architecture) :Swin Transformer 构建了一个分层的特征表示,通过逐步合并图像块(patches)来形成不同尺度的特征图,这使得模型能够适应不同尺度的图像,并具有线性的计算复杂度。
# 2. 滑动窗口机制 (Shifted Windows) :Swin Transformer在局部不重叠的窗口内计算自注意力,并允许跨窗口连接。这种机制通过限制自注意力计算在局部窗口内,同时允许窗口间的连接,提高了效率。
# 3. 线性计算复杂度 :与图像大小相关的线性计算复杂度使得 Swin Transformer 能够处理高分辨率的图像,同时保持较低的计算成本。
# 4. 多尺度特征提取 :通过特征融合的方式(Patch Merging), Swin Transformer 能够在每次特征抽取后进行下采样,增加了对原始图像的多尺度特征提取能力。
# 5. 灵活性 : Swin Transformer 的设计使其能够灵活地适应各种视觉任务,并且在多个基准测试中取得了优异的性能。
# 6. 技术创新 : Swin Transformer 解决了将 Transformer 从语言领域迁移到视觉领域时面临的挑战,如视觉实体规模的巨大差异和图像像素的高分辨率。
# 7. 应用广泛 : Swin Transformer 在图像分类、目标检测、语义分割等多个视觉任务中都取得了显著的性能提升,并有潜力取代传统的CNN架构。
# Swin Transformer 的这些特性使其成为了计算机视觉领域的一个重要里程碑,展示了基于 Transformer 模型在视觉任务中的潜力和有效性。
# 这段代码定义了一个名为 WindowAttention 的 PyTorch 模块类,它是 nn.Module 的子类。 WindowAttention 类旨在实现 Transformer 架构中的窗口化(window-based)注意力机制,这在处理图像或序列数据时特别有用,因为它可以减少全局注意力计算的复杂度。
class WindowAttention(nn.Module):
# 1. dim :表示输入特征的维度。
# 2. window_size :表示窗口的尺寸,通常是一个元组或列表,包含高度和宽度两个维度的信息(例如, (Wh, Ww) )。
# 3. num_heads :表示多头注意力机制中的头数。
# 4. qkv_bias :一个布尔值,指示是否在线性层(Q、K、V)中添加偏置项,默认为 True 。
# 5. qk_scale :一个可选的标量值,用于缩放查询(Q)和键(K)之间的点积结果。如果不提供,则会自动计算。
# 6. attn_drop :注意力权重的 dropout 比率,默认为 0. ,表示不应用 dropout。
# 7. proj_drop :输出特征的 dropout 比率,默认为 0. ,表示不应用 dropout。
def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
# 这行代码调用了父类 nn.Module 的构造函数,是初始化 PyTorch 模块时的标准做法。
super().__init__()
# 将传入的 dim 参数值赋给类的 dim 属性,这个属性表示输入特征的维度。
self.dim = dim
# 将传入的 window_size 参数值赋给类的 window_size 属性,这个属性表示窗口的尺寸,通常是一个元组或列表,包含高度和宽度两个维度的信息(例如, (Wh, Ww) )。
self.window_size = window_size # Wh, Ww
# 将传入的 num_heads 参数值赋给类的 num_heads 属性,这个属性表示多头注意力机制中的头数。
self.num_heads = num_heads
# 计算每个头的维度。这是通过将输入特征的维度 dim 除以头数 num_heads 得到的。
head_dim = dim // num_heads
# 计算缩放因子。如果提供了 qk_scale 参数,则使用它;如果没有提供,则使用每个头的维度 head_dim 的负平方根作为缩放因子。这个缩放因子用于调整查询(Q)和键(K)之间的点积结果,以防止梯度消失或爆炸。
self.scale = qk_scale or head_dim ** -0.5
# 定义相对位置偏差的参数表。
# define a parameter table of relative position bias
# 初始化了一个名为 self.relative_position_bias_table 的参数,它是一个 nn.Parameter 对象,用于存储 相对位置偏置 。这个参数是 Transformer 中相对位置编码的关键部分,它允许模型捕捉序列中不同位置之间的关系。
# (2 * window_size[0] - 1) * (2 * window_size[1] - 1) :这个维度表示相对位置偏置的数量,它是基于窗口大小计算的。
# window_size[0] 和 window_size[1] 分别是窗口的高度和宽度,通过将它们各自扩展到两倍大小再减去1,可以得到窗口内所有可能的相对位置偏置的数量。
# num_heads :这个维度表示多头注意力机制中的头数,每个头都有自己的位置偏置。
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH
# 相对位置偏置表的引入是为了在自注意力机制中考虑序列中元素的相对位置信息。在窗口化的注意力中,这个偏置表会被用来调整每个窗口内元素的注意力分数,从而使模型能够更好地理解局部结构和模式。这些偏置值将在训练过程中被优化,以帮助模型学习最有效的位置编码策略。
# torch.arange(start=0, end=0, step=1, out=None, dtype=None, device=None, requires_grad=False) -> Tensor
# torch.arange 是 PyTorch 中的一个函数,它返回一个由连续整数组成的一维张量,类似于 Python 的内置函数 range 。这个函数通常用于创建序列或索引。
# 参数说明 :
# start :序列的起始值,默认为 0。
# end :序列的结束值,但不包含此值。
# step :序列中每个元素的步长,默认为 1。
# out :一个可选的 Tensor ,用于存储输出结果。
# dtype :输出张量的所需数据类型,默认为 torch.float32 。
# device :输出张量的所需设备(CPU 或 GPU),默认为 CPU。
# requires_grad :是否需要计算梯度,默认为 False。
# 返回值 :
# 返回一个一维 Tensor ,包含了从 start 到 end (不包含)的整数序列,步长为 step 。
# 获取窗口内每个标记的成对相对位置索引。
# get pair-wise relative position index for each token inside the window
# 创建一个从0到 window_size[0]-1 的一维张量,表示窗口在高度方向上的坐标。
coords_h = torch.arange(self.window_size[0])
# 创建一个从0到 window_size[1]-1 的一维张量,表示窗口在宽度方向上的坐标。
coords_w = torch.arange(self.window_size[1])
# torch.meshgrid(*tensors, indexing='xy')
# torch.meshgrid 是 PyTorch 中的一个函数,它用于生成坐标网格,这在多维空间中可视化或计算点的坐标时非常有用。
# 参数说明 :
# tensors :一个或多个张量,这些张量定义了网格的维度。
# indexing :一个字符串,指定索引的顺序。默认为 'xy',表示笛卡尔坐标系中的顺序。也可以设置为 'ij',表示矩阵索引顺序。
# 返回值 :
# 返回一组张量,每个张量代表网格的一个维度的坐标。
# torch.stack(tensors, dim=0, out=None) → Tensor
# torch.stack 函数用于将一系列张量沿着一个新的维度连接起来。与 torch.cat (concatenate)不同, torch.stack 会创建一个新的维度来存放输入张量,而不是在已有的维度上进行连接。
# tensors :一个序列(如元组或列表)的张量,它们需要有相同的形状。
# dim :沿着哪个维度堆叠张量。默认为0,表示在最前面添加一个新的维度。
# out :一个可选的张量,用于存储输出结果。
# 返回值:
# 返回一个新的张量,它是输入张量沿着 dim 维度堆叠的结果。
# 使用 torch.meshgrid 生成高度和宽度坐标的网格,然后使用 torch.stack 将它们堆叠成一个二维张量,形状为 (2, Wh, Ww) ,其中第一个维度表示高度和宽度。
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
# torch.flatten(input, start_dim=0, end_dim=-1)
# 在 PyTorch 中, flatten 函数用于将多维张量(tensor)展平成一维张量。这个函数可以指定从哪个维度开始展平,直到张量的末尾。
# 参数说明 :
# input :要展平的输入张量。
# start_dim=0 :展平操作开始的维度。默认从第0维(即第一个维度)开始展平。
# end_dim=-1 :展平操作结束的维度。默认到张量的最后一个维度结束。如果设置为 -1 ,则表示到最后一个维度。
# 返回值 :
# 返回一个新的展平后的张量。
# 将 coords 张量在第二个维度上展平,得到形状为 (2, Wh*Ww) 的张量。
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
# 计算相对坐标,通过将展平的坐标张量自身相减得到,形状为 (2, Wh*Ww, Wh*Ww) 。
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
# 重新排列相对坐标张量的维度,并使其在内存中连续,得到形状为 (Wh*Ww, Wh*Ww, 2) 的张量。
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
# 将相对坐标的第一个维度(高度)偏移,使其从0开始。
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
# 将相对坐标的第二个维度(宽度)偏移,使其从0开始。
relative_coords[:, :, 1] += self.window_size[1] - 1
# 将相对坐标的第一个维度(高度)乘以宽度的扩展因子,以考虑宽度方向上的位置。
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
# 计算相对位置的索引,通过在最后一个维度上求和得到,形状为 (Wh*Ww, Wh*Ww) 。
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
# torch.nn.Module.register_buffer(self, name, tensor)
# 在 PyTorch 中, register_buffer 是 nn.Module 类的一个方法,用于注册一个缓冲区(buffer)到模块中。缓冲区是模块状态的一部分,但它不会被优化器视为参数,也就是说,在模型训练过程中,缓冲区的内容不会被梯度更新。这通常用于存储不需要梯度的中间结果或者固定的参数。
# 参数说明 :
# self :当前的 nn.Module 实例。
# name :缓冲区的名称,这个名称必须是唯一的,并且不能与模块中的任何参数或子模块的名称冲突。
# tensor :要注册的缓冲区张量。
# 作用 :
# register_buffer 方法允许你将一个张量注册为模块的缓冲区,这样在保存和加载模型时,这个缓冲区也会被保存和加载。这对于需要在模型的不同部分之间共享数据,或者在模型的不同阶段使用固定数据的情况非常有用。
# 将 relative_position_index 张量注册为类的缓冲区,这意味着它将被保存在模型的状态中,并且在模型的参数字典中可见,但它不会被优化器更新。
self.register_buffer("relative_position_index", relative_position_index)
# 创建一个 nn.Linear 线性层,用于生成查询(Q)、键(K)和值(V)。输入维度是 dim ,输出维度是 dim * 3 ,因为需要同时输出三个向量。 bias 参数根据 qkv_bias 决定是否在线性层中包含偏置项。
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
# 创建一个 nn.Dropout 层,用于在注意力计算中应用dropout,以减少过拟合。 attn_drop 参数指定了dropout的概率。
self.attn_drop = nn.Dropout(attn_drop)
# 创建另一个 nn.Linear 线性层,用于将注意力的输出投影回原始维度。这里输入和输出维度都是 dim 。
self.proj = nn.Linear(dim, dim)
# 创建另一个 nn.Dropout 层,用于在投影后的应用 dropout 。 proj_drop 参数指定了dropout的概率。
self.proj_drop = nn.Dropout(proj_drop)
# torch.nn.init.normal_(tensor, mean=0., std=1.)
# nn.init.normal_ 是 PyTorch 中的一个函数,用于将正态分布的值初始化到张量中。这个函数会就地修改张量的值,即直接在原始张量上进行操作,而不是返回一个新的张量。
# 参数说明 :
# tensor :要初始化的张量。
# mean=0. :正态分布的均值,默认为0。
# std=1. :正态分布的标准差,默认为1。
# 作用 :
# nn.init.normal_ 函数使用指定的均值和标准差,将正态分布的值填充到张量中。这种初始化方法可以帮助神经网络在训练初期避免梯度消失或爆炸的问题,因为它为权重提供了一个合适的初始范围。
# nn.init.normal_ 是权重初始化的常用方法之一,特别是在需要随机但有特定分布的初始值时。这种初始化方法有助于神经网络的收敛和性能。
# 使用正态分布初始化相对位置偏置表 self.relative_position_bias_table ,标准差为 0.02 。这是权重初始化的一种常见方法,有助于在训练开始时打破对称性。
nn.init.normal_(self.relative_position_bias_table, std=.02)
# 创建一个 nn.Softmax 层,用于在计算注意力分数后进行归一化。 dim=-1 指定了在最后一个维度上应用softmax函数,这通常是特征维度。
self.softmax = nn.Softmax(dim=-1)
# WindowAttention 类的 forward 方法,它实现了窗口化多头自注意力机制的前向传播逻辑。
# 接受输入张量 x 和一个可选的掩码 mask 。
def forward(self, x, mask=None):
# 解构输入张量 x 的形状,其中 B_ 是批次大小, N 是序列长度(或窗口内元素的数量), C 是特征维度。
B_, N, C = x.shape
# 调用 self.qkv 线性层处理输入 x ,得到查询(Q)、键(K)和值(V)。
# 将输出重塑为 (B_, N, 3, self.num_heads, C // self.num_heads) 的形状,其中 3 表示 Q、K 和 V 的组合。
# 使用 permute 方法重新排列维度,以便于后续的操作。 2, 0, 3, 1, 4 表示新的维度顺序。
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
# 将 qkv 张量分解为查询 q 、键 k 和值 v 。这样做是为了确保 TorchScript 的兼容性,因为 TorchScript 不支持将张量作为元组直接解构。
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) 使 torchscript 正常运行(不能使用张量作为元组)。
# 将查询 q 乘以缩放因子 self.scale ,这有助于控制梯度的规模。
q = q * self.scale
# 计算查询 q 和键 k 的转置之间的点积,得到注意力分数。 @ 表示矩阵乘法, k.transpose(-2, -1) 将键 k 的最后两个维度交换,以匹配矩阵乘法的要求。
attn = (q @ k.transpose(-2, -1))
# 使用 self.relative_position_bias_table (相对位置偏置表)和 self.relative_position_index (相对位置索引)来索引并获取相应的相对位置偏置。
# self.relative_position_index.view(-1) 将相对位置索引展平成一维张量。
# 获取的相对位置偏置被重塑为 (self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], num_heads) 的形状,其中 num_heads 是头的数量。
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
# 使用 permute 方法重新排列相对位置偏置张量的维度,使其形状变为 (num_heads, self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1]) 。
# contiguous() 方法确保张量在内存中是连续的,这对于某些操作(如维度置换)是必要的。
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
# relative_position_bias.unsqueeze(0) 在相对位置偏置张量的第一个维度(批次维度)上增加一个维度,使其形状变为 (1, num_heads, self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1]) 。
# 将这个偏置加到注意力分数 attn 上,以纳入相对位置信息。
attn = attn + relative_position_bias.unsqueeze(0)
# 目的是将相对位置偏置添加到注意力分数中,从而使模型能够捕捉序列中元素的相对位置信息。这种机制对于处理像图像或序列这样的有序数据特别重要,因为它允许模型区分不同位置之间的关系,即使这些位置在绝对值上相等。通过这种方式,模型可以更好地理解局部结构和模式。
# 处理掩码(Mask)。检查是否提供了掩码 mask 。
if mask is not None:
# 如果提供了掩码,获取掩码的形状,并从中提取第一个维度 nW ,这通常代表窗口的数量。
nW = mask.shape[0]
# 将注意力分数 attn 重塑为 (B_ // nW, nW, self.num_heads, N, N) 的形状,以匹配掩码的形状。
# 将掩码 mask 在第1维和第0维上增加一个维度,使其形状变为 (1, 1, self.num_heads, N, N) 。
# 将掩码加到重塑后的注意力分数上,以屏蔽(通常是将其设置为一个非常大的负数)不需要关注的位置。
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
# 将注意力分数和掩码相加后的结果重塑回 (-1, self.num_heads, N, N) 的形状,以便于应用softmax。
attn = attn.view(-1, self.num_heads, N, N)
# 应用 self.softmax 层对注意力分数进行归一化,使其在每个头的每一行上都和为1。
attn = self.softmax(attn)
else:
# 如果没有提供掩码,则直接对注意力分数应用softmax归一化。
attn = self.softmax(attn)
# 对归一化后的注意力分数应用 self.attn_drop ,dropout层,以减少过拟合。
attn = self.attn_drop(attn)
# 这段代码的目的是确保注意力机制能够正确地处理输入序列中的掩码(例如,用于处理序列中的填充元素),并通过softmax归一化和dropout来调整注意力分数。
# 掩码的使用在处理变长序列或填充序列时尤为重要,因为它允许模型忽略这些不需要关注的位置。
# 通过这种方式,模型可以更加专注于序列中有效的部分。
# print(attn.dtype, v.dtype)
# 注意力加权计算。
# 开始一个try块,用于捕获可能发生的异常。
try:
# 使用矩阵乘法( @ )将注意力分数 attn 与值 v 相乘,得到加权的值。
# transpose(1, 2) 将结果的第二和第三维度交换,以匹配原始输入的形状。
# reshape(B_, N, C) 将结果重塑为原始输入张量的形状,其中 B_ 是批次大小, N 是序列长度, C 是特征维度。
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
# 异常处理。如果尝试块中的操作引发异常,执行这个异常处理块。
except:
#print(attn.dtype, v.dtype)
# 如果发生异常,可能是因为张量的数据类型不匹配或设备不支持。这里, attn.half() 将注意力分数转换为半精度浮点数(FP16),然后再次尝试执行矩阵乘法和后续操作。
x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
# 线性投影和Dropout。
# 将加权的值通过 self.proj 线性层进行投影,以将特征维度映射回原始维度。
x = self.proj(x)
# 对投影后的输出应用 self.proj_drop dropout层,以减少过拟合。
x = self.proj_drop(x)
# 返回最终的输出张量 x 。
return x
# 这段代码定义了一个名为 Mlp 的类,它是 PyTorch 的 nn.Module 的子类。 Mlp 类实现了一个多层感知机(MLP)结构,通常用于深度学习模型中的全连接层。
class Mlp(nn.Module):
# 1.in_features :输入特征的维度。
# 2.hidden_features=None :隐藏层的维度。如果未指定(即为 None ),则默认与 in_features 相同。
# 3.out_features=None :输出特征的维度。如果未指定(即为 None ),则默认与 in_features 相同。
# 4.act_layer=nn.SiLU :激活函数层的类型。默认使用 nn.SiLU (也称为 Swish 激活函数),但可以替换为其他类型的激活函数。
# 5.drop=0. :Dropout 层的概率。默认为 0. ,表示不应用 Dropout。
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.):
# 调用父类 nn.Module 的构造函数。
super().__init__()
# 这行代码检查是否提供了 out_features 参数。如果没有提供(即为 None ),则将 out_features 设置为与输入特征数 in_features 相同。
out_features = out_features or in_features
# 类似地,这行代码检查是否提供了 hidden_features 参数。如果没有提供,将其设置为与输入特征数 in_features 相同。
hidden_features = hidden_features or in_features
# 创建第一个全连接层 fc1 ,它将输入特征从 in_features 维度映射到 hidden_features 维度。
self.fc1 = nn.Linear(in_features, hidden_features)
# 创建激活函数层 act ,使用传入的 act_layer 类创建实例,默认为 nn.SiLU 。
self.act = act_layer()
# 创建第二个全连接层 fc2 ,它将隐藏特征从 hidden_features 维度映射到 out_features 维度。
self.fc2 = nn.Linear(hidden_features, out_features)
# 创建一个 Dropout 层 drop ,用于正则化和防止过拟合,dropout 概率由 drop 参数指定。
self.drop = nn.Dropout(drop)
# 这段代码的目的是完成 MLP 结构的构建,包括两个全连接层和一个激活函数层,以及一个 Dropout 层。这种结构在神经网络中常用于处理序列或特征数据,特别是在 Transformer 架构中,MLP 被用作 Feed-Forward Network(FFN)的一部分。
def forward(self, x):
# 输入张量 x 通过第一个全连接层 fc1 ,这一层将输入特征从 in_features 维度映射到 hidden_features 维度。
x = self.fc1(x)
# 经过 fc1 的输出通过激活函数层 act ,这层通常用于引入非线性,以便网络可以学习复杂的函数映射。
x = self.act(x)
# 经过激活函数的输出通过 Dropout 层 drop ,这层随机将一部分输出置零,以减少过拟合。
x = self.drop(x)
# 经过 Dropout 的输出通过第二个全连接层 fc2 ,这一层将特征从 hidden_features 维度映射回 out_features 维度。
x = self.fc2(x)
# 经过 fc2 的输出再次通过 Dropout 层 drop ,进一步减少过拟合。
x = self.drop(x)
# 返回最终的输出张量 x 。
return x
# 这个方法简洁地展示了数据在 MLP 中的流动过程,包括两次全连接变换、一次非线性激活以及两次 Dropout 正则化。这种结构使得 MLP 能够学习输入数据的复杂表示,并在各种任务中发挥作用,如分类、回归或特征提取。
# 这段代码定义了一个名为 window_partition 的函数,它用于将输入的四维特征图 x 分割成多个窗口。这种操作在 Vision Transformer(ViT)和相关变体中很常见,用于实现窗口化的自注意力机制。
# 1.x :是输入的特征图。
# 2.window_size :是每个窗口的大小。
def window_partition(x, window_size):
# 解构输入特征图 x 的形状,其中 B 是批次大小, H 是特征图的高度, W 是特征图的宽度, C 是通道数。
B, H, W, C = x.shape
# 断言语句,确保特征图的高度和宽度都能被窗口大小整除。如果不能整除,将抛出一个异常。
assert H % window_size == 0, 'feature map h and w can not divide by window size' # 特征图 h 和 w 不能除以窗口大小。
# 将输入特征图 x 重塑为一个新的形状,其中高度和宽度被分割成多个 window_size x window_size 的窗口。
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
# permute(0, 1, 3, 2, 4, 5) :重新排列张量 x 的维度,以将窗口的维度放在一起。
# contiguous() :确保张量在内存中是连续的,这是进行视图操作前的必要步骤。
# view(-1, window_size, window_size, C) :将重新排列的张量重塑为最终的窗口形状,其中 -1 表示自动计算该维度的大小。
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
# 返回 分割后的窗口 张量。
return windows
# 这个函数的目的是将输入的特征图分割成多个较小的窗口,以便在每个窗口内独立地应用自注意力机制。这种方法可以减少自注意力计算的复杂度,特别是当处理大规模特征图时。通过这种方式,模型可以更有效地处理图像或序列数据,同时保持对局部特征的敏感性。
# 这段代码定义了一个名为 window_reverse 的函数,它用于将通过 window_partition 函数分割得到的窗口重新组合回原始的特征图。这种操作在 Vision Transformer(ViT)和相关变体中用于在应用窗口化自注意力后 恢复特征图 。
# 1.windows :是分割后的窗口张量。
# 2.window_size :是每个窗口的大小。
# 3.H :是原始特征图的高度。
# 4.W :是原始特征图的宽度。
def window_reverse(windows, window_size, H, W):
# 计算批次大小 B 。这个计算基于窗口张量的第一个维度( windows.shape[0] )和原始特征图的尺寸,确保重新组合后的张量形状正确。
B = int(windows.shape[0] / (H * W / window_size / window_size))
# 将窗口张量 windows 重塑为一个新的形状,其中高度和宽度被分割成多个 window_size x window_size 的窗口,并且每个窗口内的特征被展平。
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
# permute(0, 1, 3, 2, 4, 5) :重新排列张量 x 的维度,以将窗口的行和列维度与原始特征图的行和列维度对齐。
# contiguous() :确保张量在内存中是连续的,这是进行视图操作前的必要步骤。
# view(B, H, W, -1) :将重新排列的张量重塑为原始特征图的形状,其中 -1 表示自动计算该维度的大小,即通道数。
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
# 返回重新组合后的特征图张量。
return x
# 这个函数的目的是将分割后的窗口张量恢复为原始的特征图,以便可以继续进行后续的网络层处理或输出最终结果。这种方法允许模型在局部窗口内捕捉特征的同时,保持对全局特征的感知。通过这种方式,模型可以更有效地处理图像或序列数据,同时保持对局部和全局特征的敏感性。
# 这段代码定义了一个名为 SwinTransformerLayer 的类,它是 PyTorch 的 nn.Module 的子类。 SwinTransformerLayer 类实现了 Swin Transformer 架构中的一个单独的 Transformer 层。
class SwinTransformerLayer(nn.Module):
# 这是类的构造函数,用于初始化类的实例。
# 1.dim :输入特征的维度。
# 2.num_heads :多头注意力机制中的头数。
# 3.window_size=8 :窗口化的自注意力机制中窗口的大小,默认为8。
# 4.shift_size=0 :窗口分割时的偏移量,默认为0,这意味着不进行窗口偏移。
# 5.mlp_ratio=4. :多层感知机(MLP)的隐藏层维度与输入维度的比例,默认为4。
# 6.qkv_bias=True :在线性层(Q、K、V)中是否添加偏置项。
# 7.qk_scale=None :查询(Q)和键(K)之间的点积结果的缩放因子。
# 8.drop=0. : dropout 层的概率,默认为0,表示不应用 dropout。
# 9.attn_drop=0. :注意力权重的 dropout 概率,默认为0。
# 10.drop_path=0. : Stochastic Depth 的概率,默认为0。
# 11.act_layer=nn.SiLU :激活函数层的类型,默认使用 nn.SiLU 。
# 12.norm_layer=nn.LayerNorm :归一化层的类型,默认使用 nn.LayerNorm 。
def __init__(self, dim, num_heads, window_size=8, shift_size=0,
mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.SiLU, norm_layer=nn.LayerNorm):
super().__init__()
# 将传入的 dim 参数值赋给类的 dim 属性,这个属性表示输入特征的维度。
self.dim = dim
# 将传入的 num_heads 参数值赋给类的 num_heads 属性,这个属性表示多头注意力机制中的头数。
self.num_heads = num_heads
# 将传入的 window_size 参数值赋给类的 window_size 属性,这个属性表示窗口化的自注意力机制中窗口的大小。
self.window_size = window_size
# 将传入的 shift_size 参数值赋给类的 shift_size 属性,这个属性表示在窗口分割时的偏移量。 shift_size 用于控制窗口是否需要偏移,以及偏移的大小。
self.shift_size = shift_size
# 将传入的 mlp_ratio 参数值赋给类的 mlp_ratio 属性,这个属性表示多层感知机(MLP)的隐藏层维度与输入维度的比例。这个比例用于确定 MLP 中的隐藏层大小。
self.mlp_ratio = mlp_ratio
# if min(self.input_resolution) <= self.window_size:
# # if window size is larger than input resolution, we don't partition windows
# self.shift_size = 0
# self.window_size = min(self.input_resolution)
# 断言语句,确保 shift_size 的值在0到 window_size (不包括 window_size )之间。这是必要的,因为窗口偏移量必须小于窗口大小。
assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
# torch.nn.LayerNorm(input_shape, eps=1e-5, elementwise_affine=True)
# nn.LayerNorm 是 PyTorch 中的一个类,它实现了层归一化(Layer Normalization)。层归一化是一种归一化技术,用于提高深度神经网络的训练速度和稳定性。它与批量归一化(Batch Normalization)不同,因为它是在单个样本的特征上进行归一化,而不是在整个批次上。
# 参数说明 :
# input_shape :一个元组,表示输入张量的形状。对于 nn.LayerNorm ,输入张量应该是至少二维的,其中最后一个维度(除了批次维度)是归一化计算的特征维度。
# eps :一个很小的值,用于避免除以零的情况。默认值为 1e-5 。
# elementwise_affine :一个布尔值,指示是否包含可学习的尺度(scale)和偏移(shift)参数。默认值为 True ,表示包含。
# 返回值 :
# 返回一个 LayerNorm 模块的实例,可以应用于输入数据进行归一化。
# nn.LayerNorm 的工作原理是 :
# 对于每个样本,它计算最后一个维度上的特征的均值和方差,然后使用这些统计数据来归一化特征,使得归一化后的特征的均值为 0,方差为 1。
# 如果 elementwise_affine=True ,则归一化后的特征还会乘以一个可学习的尺度参数并加上一个可学习的偏移参数。
# 创建第一个归一化层 norm1 ,使用传入的 norm_layer 类(默认为 nn.LayerNorm )和输入特征的维度 dim 。
self.norm1 = norm_layer(dim)
# class WindowAttention(nn.Module):
# -> def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
# 创建自注意力机制 attn ,使用 WindowAttention 类,传入参数包括特征维度 dim 、窗口大小 (self.window_size, self.window_size) 、头数 num_heads 、是否使用偏置 qkv_bias 、缩放因子 qk_scale 、注意力dropout attn_drop 和投影dropout drop 。
self.attn = WindowAttention(
dim, window_size=(self.window_size, self.window_size), num_heads=num_heads,
qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
# 根据 drop_path 的值创建 DropPath 层或恒等层 nn.Identity 。如果 drop_path 大于0,则使用 DropPath 层,否则使用恒等层。DropPath 是一种正则化技术,用于随机丢弃路径以防止过拟合。
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
# 创建第二个归一化层 norm2 ,与第一个归一化层( norm1 )类似。
self.norm2 = norm_layer(dim)
# 计算 MLP 的隐藏层维度,基于输入特征维度 dim 和比例 mlp_ratio 。
mlp_hidden_dim = int(dim * mlp_ratio)
# class Mlp(nn.Module):
# -> def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.):
# 创建 MLP mlp ,使用 Mlp 类,传入参数包括 输入特征维度 dim 、隐藏层维度 mlp_hidden_dim 、激活函数层 act_layer 和 dropout 概率 drop 。
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
# 用于创建一个注意力掩码(mask),用于空间窗口注意力机制(Spatial Window Attention with Multi-Scale Attention,简称SW-MSA)。这个掩码用于指示在计算自注意力时哪些位置是有效的。
# 1.H :图像的高度。
# 2.W :图像的宽度。
def create_mask(self, H, W):
# 计算 SW-MSA 的注意力掩码。
# calculate attention mask for SW-MSA
# 创建了一个形状为 (1, H, W, 1) 的张量,初始值都为0。这个张量将被用来存储掩码值。
img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1
# self.window_size 和 self.shift_size 是模型中的参数,分别代表 窗口的大小 和 窗口之间的偏移量 。
# h_slices 和 w_slices : 这两个元组分别用于处理高度和宽度方向上的切片。
# 每个元组包含三个切片对象,对应于不同的窗口区域 :
# 第一个切片 ( slice(0, -self.window_size) ) :表示从当前窗口的开始到窗口大小之前的位置。
# 第二个切片 ( slice(-self.window_size, -self.shift_size) ) :表示从窗口大小之前的位置到位移大小之前的位置。
# 第三个切片 ( slice(-self.shift_size, None) ) :表示从位移大小之前的位置到窗口的结束。
h_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
# 初始化一个计数器,用于为每个窗口分配一个唯一的ID。
cnt = 0
# h_slices 和 w_slices 每个都包含三个切片对象,分别对应于一个 窗口的开始部分 、 中间偏移部分 和 结束部分 。
# 循环遍历高度方向上的每个切片。
for h in h_slices:
# 在每个高度切片内部,循环遍历宽度方向上的每个切片。
for w in w_slices:
# 对于每一对高度和宽度的切片组合,代码执行以下操作。
# img_mask[:, h, w, :] = cnt :将 img_mask 张量中对应高度切片 h 和宽度切片 w 的位置设置为当前的计数器 cnt 的值。这意味着在 img_mask 中,所有属于同一个窗口的像素点都会被标记为相同的计数器值。
img_mask[:, h, w, :] = cnt
# 在设置完当前窗口的掩码值后,计数器 cnt 增加1,为下一个窗口准备一个新的唯一标识。
cnt += 1
# 这样,每个窗口都会被分配一个唯一的ID,这些ID在自注意力机制中用来区分不同的窗口。这种方法允许模型在计算自注意力时只考虑窗口内部的像素,而忽略窗口外部的像素,从而减少计算量并提高效率。
# def window_partition(x, window_size): -> 它用于将输入的四维特征图 x 分割成多个窗口。返回 分割后的窗口 张量。 -> return windows
# 这行代码将之前创建的 img_mask 张量分割成多个窗口。 window_partition 函数将 img_mask 按照 self.window_size 划分成多个小窗口,每个窗口的大小为 window_size x window_size 。
# 结果是一个形状为 (nW, window_size, window_size, 1) 的张量,其中 nW 是窗口的总数。
mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1
# 这行代码将 mask_windows 张量重塑(view)为一个二维张量,形状为 (-1, self.window_size * self.window_size) 。这里 -1 表示自动计算该维度的大小,以确保总元素数量保持不变。
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
# 这行代码计算了自注意力掩码。
# 首先, unsqueeze 操作在指定维度上增加一个维度。 unsqueeze(1) 和 unsqueeze(2) 分别在第二个和第三个维度上增加维度,使得张量的形状变为 (-1, 1, self.window_size * self.window_size) 和 (-1, self.window_size * self.window_size, 1) 。
# 然后,通过相减操作,对于每个窗口内的元素,如果它们属于同一个窗口(即它们的 mask_windows 值相同),则相减结果为0;如果它们属于不同窗口,则相减结果非0。这样,我们可以得到一个表示元素间是否属于同一窗口的掩码。
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
# masked_fill(mask, value)
# masked_fill_(mask, value)
# masked_fill 是 PyTorch 中的一个函数,它用于根据给定的掩码(mask)将张量(tensor)中指定位置的元素替换为特定的值。
# 参数 :
# mask :一个布尔类型的张量(BoolTensor),其形状必须与要操作的张量形状相同或者可以广播到张量的形状。在 mask 中为 True 的位置将被替换,为 False 的位置保持不变。
# value :一个标量或与被操作的张量具有相同形状的张量,用于指定替换的值。
# 功能描述 :
# masked_fill 函数会检查 mask 中的每个元素,如果 mask 对应的位置为 True ,则在输出张量中用 value 替换原张量中对应位置的值;如果为 False ,则保留原张量中对应位置的值。
# 返回值 :
# 返回一个新的张量,其中包含了在 mask 指定位置上被替换的值,其余位置保持原样。原始张量不会被改变( masked_fill 是非原地操作)。
# 这行代码进一步处理 attn_mask ,使用 masked_fill 操作将非0值替换为一个非常大的负数( -100.0 ),这在自注意力计算中通常用于抑制(mask out)不同窗口间的相互作用。
# 同时,将0值替换为0.0,表示同一窗口内的元素可以相互影响。
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
# 最后,返回处理后的注意力掩码 attn_mask ,它将被用于自注意力机制中,以确保模型在计算自注意力时只考虑窗口内部的相互作用。
return attn_mask
# 总结来说,这段代码通过分割图像掩码、重塑、相减和掩码填充等操作,生成了一个用于自注意力机制的注意力掩码,该掩码能够区分窗口内和窗口外的元素,从而实现有效的局部特征提取和计算效率提升。
def forward(self, x):
# 这是一个注释,说明接下来的代码将把输入张量 x 的形状从 [batch_size, channels, height, width] 重塑为 [batch_size, length, channels] ,其中 length 是窗口数量乘以每个窗口的元素数量。
# reshape x[b c h w] to x[b l c]
# 这行代码获取输入张量 x 的形状,并分别赋值给 _ (忽略的通道数)、 H_ (高度)和 W_ (宽度)。
_, _, H_, W_ = x.shape
# 初始化一个变量 Padding ,用于标记是否需要对输入张量进行填充。
Padding = False
# 这行代码检查输入张量的高度和宽度是否小于窗口大小,或者是否能被窗口大小整除。如果不能,就需要进行填充。
if min(H_, W_) < self.window_size or H_ % self.window_size!=0 or W_ % self.window_size!=0:
# 如果需要填充,将 Padding 设置为 True 。
Padding = True
# print(f'img_size {min(H_, W_)} is less than (or not divided by) window_size {self.window_size}, Padding.')
# 计算需要在 宽度 和 高度 方向上分别填充多少才能达到窗口大小的倍数。这里使用了 模运算 和 取余数 的操作。
pad_r = (self.window_size - W_ % self.window_size) % self.window_size
pad_b = (self.window_size - H_ % self.window_size) % self.window_size
# 使用PyTorch的 F.pad 函数对输入张量 x 进行填充。 F.pad 的参数是一个元组,指定了在每个方向上需要填充的像素数。这里 (0, pad_r, 0, pad_b) 表示在宽度的右侧填充 pad_r 个像素,在高度的底部填充 pad_b 个像素。
x = F.pad(x, (0, pad_r, 0, pad_b))
# print('2', x.shape)
# 获取输入张量 x 的形状,分别赋值给 B (批量大小)、 C (通道数)、 H (高度)和 W (宽度)。
B, C, H, W = x.shape
# 计算线性化后的长度 L ,即图像的总像素数。
L = H * W
# permute(0, 2, 3, 1) :重新排列输入张量 x 的维度,将通道维度 C 移动到最后。
# contiguous() :确保张量在内存中是连续存储的,这对于某些操作(如 view )是必要的。
# view(B, L, C) :将张量 x 重塑为形状 [B, L, C] ,其中 L 是图像的总像素数。
x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C) # b, L, c
# 这是一个注释,说明接下来的代码将在 forward 方法中创建注意力掩码。
# create mask from init to forward
# 检查 self.shift_size 是否大于0,这个参数通常用于确定是否需要进行窗口偏移。
if self.shift_size > 0:
# def create_mask(self, H, W):
# -> 用于创建一个注意力掩码(mask),用于空间窗口注意力机制(Spatial Window Attention with Multi-Scale Attention,简称SW-MSA)。这个掩码用于指示在计算自注意力时哪些位置是有效的。
# -> return attn_mask
# 如果 self.shift_size 大于0,调用 create_mask 方法创建注意力掩码,并将掩码移动到输入张量 x 所在的设备(例如CPU或GPU)。
attn_mask = self.create_mask(H, W).to(x.device)
else:
# 如果 self.shift_size 不大于0,设置 attn_mask 为 None ,表示不需要注意力掩码。
attn_mask = None
# 创建一个快捷连接(shortcut),将原始的 x 张量存储起来,以便后续可能的残差连接使用。
shortcut = x
# 将输入张量 x 通过第一个归一化层(例如LayerNorm或BatchNorm)。
x = self.norm1(x)
# 将归一化后的张量 x 重塑回原始的四维形状 [B, H, W, C] ,以便进行后续的窗口划分或其他操作。
x = x.view(B, H, W, C)
# cyclic shift 循环偏移。
# 检查 self.shift_size 是否大于0,这个参数通常用于确定是否需要进行窗口偏移。
if self.shift_size > 0:
# torch.roll(input, shifts, dims=None)
# torch.roll 是 PyTorch 中的一个函数,它对输入张量进行滚动操作,也就是将张量中的元素沿着指定的维度“滚动”一定数量的位置。这个操作会将指定维度上的元素按照指定的步数向前或向后移动,移动出的元素会被重新放置到另一端。
# 参数说明 :
# input :输入的 Tensor ,即要进行滚动操作的张量。
# shifts :一个整数或整数元组,指定每个维度上元素需要滚动的步数。如果是一个整数,则在所有维度上滚动相同的步数。
# dims :一个整数或整数元组,指定要滚动的维度。如果为 None ,则滚动所有维度。
# 返回值 :
# 返回一个新的 Tensor ,其形状与输入 Tensor 相同,但元素已经根据指定的步数和维度进行了滚动。
# 如果 self.shift_size 大于0,使用 torch.roll 函数对输入张量 x 进行偏移操作。这里的偏移是在高度和宽度维度上进行的,偏移量为 -self.shift_size ,这意味着窗口会向下和向右偏移。 dims=(1, 2) 指定了偏移的维度,即高度和宽度。
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
else:
# 如果 self.shift_size 不大于0,设置 shifted_x 等于 x ,表示不需要进行窗口偏移。
shifted_x = x
# partition windows
# def window_partition(x, window_size): -> 它用于将输入的四维特征图 x 分割成多个窗口。返回 分割后的窗口 张量。 -> return windows
# 对偏移后的张量 shifted_x 进行窗口划分,将其分割成多个窗口。 window_partition 函数将张量划分为 nW*B, window_size, window_size, C 的形状,其中 nW 是窗口的数量。
x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C
# 将窗口划分后的张量 x_windows 重塑为 nW*B, window_size*window_size, C 的形状,以适应自注意力机制的输入要求。
x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
# W-MSA/SW-MSA
# 应用自注意力机制(W-MSA/SW-MSA),将重塑后的窗口张量 x_windows 和注意力掩码 attn_mask 作为输入。输出 attn_windows 的形状为 nW*B, window_size*window_size, C 。
attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C
# merge windows
# 将自注意力机制的输出 attn_windows 重塑回 -1, self.window_size, self.window_size, C 的形状,以准备进行窗口合并。
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
# def window_reverse(windows, window_size, H, W):
# -> 它用于将通过 window_partition 函数分割得到的窗口重新组合回原始的特征图。这种操作在 Vision Transformer(ViT)和相关变体中用于在应用窗口化自注意力后 恢复特征图 。返回重新组合后的特征图张量。
# -> return x
# 使用 window_reverse 函数将自注意力机制的输出 attn_windows 合并回原始图像的形状 B, H', W', C 。这里的 H' 和 W' 可能是经过填充后的高度和宽度。
shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C
# reverse cyclic shift
# 检查 self.shift_size 是否大于0,这个参数通常用于确定是否需要进行窗口偏移的反向操作。
if self.shift_size > 0:
# 如果 self.shift_size 大于0,使用 torch.roll 函数对偏移后的张量 shifted_x 进行反向偏移操作。这里的偏移是在高度和宽度维度上进行的,偏移量为 self.shift_size ,这意味着窗口会向上和向左回滚。
x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
# 如果 self.shift_size 不大于0,设置 x 等于 shifted_x ,表示不需要进行窗口偏移的反向操作。
x = shifted_x
# 将反向偏移后的张量 x 重塑为形状 [B, H * W, C] ,其中 H * W 是图像的总像素数。
x = x.view(B, H * W, C)
# 这是一个注释,说明接下来的代码将通过一个前馈网络(FFN)。
# FFN
# 将前馈网络的输出与一个快捷连接(shortcut)相加, shortcut 通常是原始输入或者经过某些操作的原始输入。 self.drop_path 是一个随机丢弃路径的函数,用于正则化。
x = shortcut + self.drop_path(x)
# 将经过第二个归一化层 self.norm2 和前馈网络 self.mlp 处理后的输出与之前的 x 相加,形成残差连接。 self.drop_path 同样用于正则化。
x = x + self.drop_path(self.mlp(self.norm2(x)))
# 重新排列张量 x 的维度,将通道维度 C 移动到高度和宽度之前。
# contiguous() 确保张量在内存中是连续存储的。
# view(-1, C, H, W) 将张量重塑为形状 [-1, C, H, W] ,恢复到原始的四维形状。
x = x.permute(0, 2, 1).contiguous().view(-1, C, H, W) # b c h w
# 检查是否进行了填充操作。
if Padding:
# 如果进行了填充,将张量 x 裁剪回原始的高度 H_ 和宽度 W_ ,以撤销之前进行的填充操作。
x = x[:, :, :H_, :W_] # reverse padding
# 返回最终的输出张量 x 。
return x
# 这段代码展示了如何在模型的前向传播过程中实现窗口化的自注意力机制和前馈网络,并将结果恢复到原始的输入尺寸。通过这种方式,模型能够有效地捕捉局部特征,并通过残差连接和快捷连接保持信息的流动。
# 这段代码定义了一个名为 SwinTransformerBlock 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了 Swin Transformer 的一个变体,它可以根据需要对输入通道进行卷积变换,然后通过一系列 Swin Transformer 层处理输入数据。
class SwinTransformerBlock(nn.Module):
# 1.c1 (int) :输入通道数。
# 2.c2 (int) :输出通道数。如果 c1 不等于 c2 ,则会创建一个卷积层来调整通道数。
# 3.num_heads (int) :在自注意力层中使用的注意力头的数量。
# 4.num_layers (int) :Swin Transformer 层的数量。
# 5.window_size (int) :用于自注意力计算的窗口大小,默认值为 8。
def __init__(self, c1, c2, num_heads, num_layers, window_size=8):
super().__init__()
# 初始化一个成员变量 self.conv ,它将用于存储一个可能的卷积层。
self.conv = None
# 检查输入通道数 c1 和输出通道数 c2 是否不同。
if c1 != c2:
# 如果 c1 和 c2 不同,创建一个卷积层 Conv ,用于将输入通道数从 c1 转换为 c2 。
self.conv = Conv(c1, c2)
# remove input_resolution 删除input_resolution。
# 创建一个 nn.Sequential 容器,它将按顺序包含多个 SwinTransformerLayer 实例。
# 列表推导式 [SwinTransformerLayer(...) for i in range(num_layers)] 创建了 num_layers 个 SwinTransformerLayer 实例。
# shift_size 参数根据层的索引 i 交替设置为 0 和 window_size // 2 ,以实现窗口偏移的效果。
self.blocks = nn.Sequential(*[SwinTransformerLayer(dim=c2, num_heads=num_heads, window_size=window_size,
shift_size=0 if (i % 2 == 0) else window_size // 2) for i in range(num_layers)])
# 前向传播 forward 方法。
def forward(self, x):
# 检查是否存在卷积层。
if self.conv is not None:
# 如果存在卷积层,将输入 x 通过该卷积层。
x = self.conv(x)
# 将输入 x 通过 nn.Sequential 容器中的所有 Swin Transformer 层。
x = self.blocks(x)
# 返回处理后的输出 x 。
return x
# SwinTransformerBlock 类定义了一个包含可选卷积层和多个 Swin Transformer 层的神经网络模块。
# 这个模块可以用于构建 Swin Transformer 架构,它通过窗口化的自注意力机制处理输入数据,并可以通过调整 num_heads 和 num_layers 参数来控制模型的复杂度。
# 这种设计允许模型捕捉局部特征,并在不同的窗口之间传递信息,从而实现有效的特征提取和表示。
# 这段代码定义了一个名为 STCSPA 的 PyTorch 神经网络模块类,它继承自 nn.Module 。 STCSPA 类代表了一个Cross Stage Partial Network (CSP)的变体,结合了Swin Transformer Block(STB)和CSP结构。
class STCSPA(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
# 1.c1 (int) :输入通道数( ch_in )。
# 2.c2 (int) :输出通道数( ch_out )。
# 3.n (int) :Swin Transformer Block 中的层数,默认为 1( number )。
# 4.shortcut (bool) :是否使用快捷连接,默认为 True ( shortcut )。
# 5.g (int) :分组卷积的组数,默认为 1( groups )。
# 6.e (float) :扩张率,用于计算隐藏层通道数,默认为 0.5( expansion )。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
# 调用父类 nn.Module 的初始化方法。
super(STCSPA, self).__init__()
# 计算隐藏通道数 c_ ,它是输出通道数 c2 的 e 倍。
c_ = int(c2 * e) # hidden channels
# 创建三个卷积层, self.cv1 和 self.cv2 用于将输入通道 c1 转换为隐藏通道 c_ , self.cv3 用于将两个 c_ 通道的张量合并回输出通道 c2 。
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
# 计算用于Swin Transformer Block的注意力头数,这里简单地将隐藏通道数 c_ 除以 32。
num_heads = c_ // 32
# 创建一个 SwinTransformerBlock 实例,用于处理通过 self.cv1 的输出。这里 SwinTransformerBlock 的输入和输出通道数都是 c_ ,注意力头数为 num_heads ,层数为 n 。
self.m = SwinTransformerBlock(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
# 前向传播 forward 方法
def forward(self, x):
# 将输入 x 通过 self.cv1 卷积层,然后通过 SwinTransformerBlock 实例 self.m 。
y1 = self.m(self.cv1(x))
# 将输入 x 直接通过 self.cv2 卷积层。
y2 = self.cv2(x)
# 将 y1 和 y2 在通道维度( dim=1 )上进行拼接,然后通过 self.cv3 卷积层输出最终结果。
return self.cv3(torch.cat((y1, y2), dim=1))
# STCSPA 类结合了 CSP 结构和 Swin Transformer 的特点,通过将输入分成两个路径,一个路径通过Swin Transformer Block处理,另一个路径直接通过卷积层,然后将两个路径的输出合并。
# 这种设计旨在结合 CSP 的效率和 Swin Transformer 的强大特征提取能力。通过这种方式, STCSPA 可以有效地处理图像数据,并在保持计算效率的同时提高模型的性能。
# 这段代码定义了一个名为 STCSPB 的 PyTorch 神经网络模块类,它继承自 nn.Module 。 STCSPB 类代表了一个结合了 Cross Stage Partial Networks (CSP) 结构和 Swin Transformer Block 的网络模块。
class STCSPB(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(STCSPB, self).__init__()
# 计算隐藏通道数 c_ ,这里直接使用输出通道数 c2 。
c_ = int(c2) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
num_heads = c_ // 32
self.m = SwinTransformerBlock(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
# 将输入 x 通过 self.cv1 卷积层。
x1 = self.cv1(x)
# 将 x1 的输出通过 Swin Transformer Block self.m 。
y1 = self.m(x1)
# 将 x1 的输出直接通过 self.cv2 卷积层。
y2 = self.cv2(x1)
# 将 y1 和 y2 在通道维度( dim=1 )上进行拼接,然后通过 self.cv3 卷积层输出最终结果。
return self.cv3(torch.cat((y1, y2), dim=1))
# STCSPB 类是一个结合了 CSP 结构和 Swin Transformer 的网络模块,旨在提高网络的特征提取能力和计算效率。通过使用扩张卷积和 Swin Transformer Block,这个类能够在保持计算效率的同时,捕获更丰富的特征表示。
# 这个类的设计灵感来源于 WongKinYiu 的 CrossStagePartialNetworks 项目,该项目在 GitHub 上提供了更多关于 CSP 结构的实现细节和应用示例。
# 这段代码定义了一个名为 STCSPC 的 PyTorch 神经网络模块类,它继承自 nn.Module 。 STCSPC 类代表了一个结合了 Cross Stage Partial Networks (CSP) 结构和 Swin Transformer Block 的网络模块。
class STCSPC(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(STCSPC, self).__init__()
c_ = int(c2 * e) # hidden channels
# 创建四个卷积层 :
# self.cv1 :将输入通道 c1 转换为隐藏通道 c_ 。
# self.cv2 :将输入通道 c1 转换为隐藏通道 c_ ,用于快捷连接。
# self.cv3 :将 Swin Transformer Block 的输出再次转换为隐藏通道 c_ 。
# self.cv4 :将两个 c_ 通道的张量合并回输出通道 c2 。
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(c_, c_, 1, 1)
self.cv4 = Conv(2 * c_, c2, 1, 1)
num_heads = c_ // 32
self.m = SwinTransformerBlock(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
# 将输入 x 通过 self.cv1 卷积层,然后通过 Swin Transformer Block self.m ,最后通过 self.cv3 卷积层。
y1 = self.cv3(self.m(self.cv1(x)))
# 将输入 x 直接通过 self.cv2 卷积层,实现快捷连接。
y2 = self.cv2(x)
# 将 y1 和 y2 在通道维度( dim=1 )上进行拼接,然后通过 self.cv4 卷积层输出最终结果。
return self.cv4(torch.cat((y1, y2), dim=1))
# STCSPC 类是一个结合了 CSP 结构和 Swin Transformer 的网络模块,旨在提高网络的特征提取能力和计算效率。通过使用扩张卷积和 Swin Transformer Block,这个类能够在保持计算效率的同时,捕获更丰富的特征表示。
# 这个类的设计灵感来源于 WongKinYiu 的 CrossStagePartialNetworks 项目,该项目在 GitHub 上提供了更多关于 CSP 结构的实现细节和应用示例。
# 通过这种设计, STCSPC 可以在不同的网络层之间有效地共享和融合特征,提高模型的性能。
##### end of swin transformer #####
##### swin transformer v2 #####
# Swin Transformer V2 是 Swin Transformer 的改进版本,旨在解决大规模视觉模型在训练和应用中面临的挑战。
# 以下是 Swin Transformer V2 的一些关键特点和改进 :
# 1. 更大的模型容量 : Swin Transformer V2 显著增加了模型的参数量,最高可达3 Billion(30亿),以处理更高分辨率的图像。
# 2. 后规范化技术 (Post-Normalization) :V2 将 LayerNorm 从每个 Transformer 块的前面移动到后面,即在 Attention 和 MLP 之后应用归一化。这种后规范化技术有助于稳定训练过程中的激活值,使得深层和浅层的输出幅值更加稳定。
# 3. 余弦相似度 :为了解决后规范化可能导致的相似度计算问题,V2 将点乘相似度替换为余弦相似度,这有助于保持归一化效果,使得 Attention 输出更加可控。
# 4. 对数空间连续位置偏置 :V2 引入了对数空间连续位置偏置技术,使得模型能够更准确地拟合不同窗口大小的位置信息。
# 5. 自监督学习 :为了解决有监督数据不足的问题, Swin Transformer V2 引入了自监督学习的方法,通过 Mask 掉部分 token 来训练模型,以学习更好的图像表征能力。
# 6. 跨尺度训练 :V2 通过一系列技术改进,使得模型能够适应不同分辨率的图像,包括低分辨率的预训练图像和高分辨率的下游任务图像。
# 7. 性能提升 : Swin Transformer V2 在多个基准数据集上取得了新记录,包括 ImageNet 分类、COCO 检测、ADE20K 语义分割以及 Kinetics-400 动作分类。
# 8. 分层表示能力: V2 保留了 Swin Transformer 的分层表示能力,通过不同层次的 Transformer 层提取不同级别的特征。
# 总的来说,Swin Transformer V2 通过一系列的改进,提高了模型的容量和分辨率,增强了模型的训练稳定性和下游任务的适应性,使得模型在多个视觉任务上取得了显著的性能提升。
# 这段代码定义了一个名为 WindowAttention_v2 的 PyTorch 神经网络模块类,它继承自 nn.Module 。这个类实现了一个窗口化的自注意力机制,其中包括了相对位置偏置和可学习的缩放因子。
class WindowAttention_v2(nn.Module):
# 1. dim (int) : 输入特征的维度。这是每个输入向量(例如,每个像素或特征点)的通道数。
# 2. window_size (tuple or list) : 窗口的大小,以 (height, width) 的形式给出。这个参数定义了自注意力计算的局部窗口的尺寸。
# 3. num_heads (int) : 注意力头的数量。这个参数决定了自注意力机制中并行处理的数量,通常用于实现多头自注意力。
# 4. qkv_bias (bool) : 是否在线性层 qkv (即查询、键、值的线性变换)中使用偏置项。默认值为 True 。
# 5. attn_drop (float) : 注意力权重的 dropout 比率。这个参数用于正则化,防止过拟合。默认值为 0. ,表示不使用 dropout。
# 6. proj_drop (float) : 投影权重的 dropout 比率。这个参数同样用于正则化,防止过拟合。默认值为 0. ,表示不使用 dropout。
# 7. pretrained_window_size (list or tuple) : 预训练时使用的窗口大小,以 [height, width] 的形式给出。这个参数用于在迁移学习或预训练模型微调时调整窗口大小。默认值为 [0, 0] ,表示不使用预训练窗口大小或者在预训练时窗口大小与当前窗口大小相同。
def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.,
pretrained_window_size=[0, 0]):
# 这行代码调用了父类 nn.Module 的构造函数,这是 Python 中类继承的常用模式,用于确保子类正确地继承了父类的所有属性和方法。
super().__init__()
# 将传入的参数 dim 赋值给类的成员变量 self.dim 。这个变量存储了输入特征的维度。
self.dim = dim
# 将传入的参数 window_size 赋值给类的成员变量 self.window_size 。这个变量存储了窗口的尺寸, window_size 应该是一个包含两个元素的序列,分别表示窗口的高度和宽度( Wh , Ww )。
self.window_size = window_size # Wh, Ww
# 将传入的参数 pretrained_window_size 赋值给类的成员变量 self.pretrained_window_size 。这个变量存储了预训练模型中使用的窗口尺寸,用于迁移学习或模型微调时调整窗口尺寸。
self.pretrained_window_size = pretrained_window_size
# 将传入的参数 num_heads 赋值给类的成员变量 self.num_heads 。这个变量存储了多头自注意力机制中头的数量。
self.num_heads = num_heads
# 创建一个可学习的参数 logit_scale ,它是一个形状为 (num_heads, 1, 1) 的张量,每个元素初始化为 log(10) 。这个参数用于缩放注意力分数的对数,允许模型在训练过程中调整缩放因子。 requires_grad=True 表示这个参数在训练过程中需要计算梯度。
# logit_scale 参数的使用是为了在计算注意力分数时提供额外的灵活性,特别是在处理不同尺度的特征或不同复杂度的任务时。通过学习这个参数,模型可以更好地调整注意力分布,从而提高性能。
self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)
# mlp to generate continuous relative position bias
# 生成连续相对位置偏置的 MLP(多层感知机)。
# 定义了一个序列化的神经网络模块,用于生成连续的相对位置偏置。
# nn.Linear(2, 512, bias=True) :第一个全连接层,输入维度为2(代表相对位置的高和宽坐标),输出维度为512,启用偏置项。
# nn.ReLU(inplace=True) :激活函数,使用 ReLU, inplace=True 表示在原地进行计算,减少内存消耗。
# nn.Linear(512, num_heads, bias=False) :第二个全连接层,将512维的输出映射到与注意力头数相同的维度,不启用偏置项。
self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
nn.ReLU(inplace=True),
nn.Linear(512, num_heads, bias=False))
# get relative_coords_table
# 生成相对坐标表。
# relative_coords_h 和 relative_coords_w :
# 使用 torch.arange 生成两个张量,分别表示窗口内 相对位置的高和宽坐标 。坐标范围从 -(window_size[0] - 1) 到 window_size[0] (高度方向),以及从 -(window_size[1] - 1) 到 window_size[1] (宽度方向)。
relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
# relative_coords_table :
# 使用 torch.meshgrid 生成高和宽坐标的网格,然后使用 torch.stack 将这两个网格堆叠起来。
# permute(1, 2, 0) :调整张量的维度顺序,以适应后续操作。
# contiguous() :确保张量在内存中是连续的,这是某些 PyTorch 操作的要求。
# unsqueeze(0) :在张量的第一个维度上增加一个维度,使其形状变为 [1, 2*Wh-1, 2*Ww-1, 2] ,其中 Wh 和 Ww 分别是窗口的高度和宽度。
relative_coords_table = torch.stack(
torch.meshgrid([relative_coords_h,
relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0) # 1, 2*Wh-1, 2*Ww-1, 2
# 条件判断。检查预训练窗口大小是否有效(即非零)。
if pretrained_window_size[0] > 0:
# 预训练窗口大小存在时的归一化。
# 如果预训练窗口大小存在,使用它来归一化 relative_coords_table 中的坐标。
# 将高度坐标除以预训练窗口高度减一。
relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)
# 将宽度坐标除以预训练窗口宽度减一。
relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)
else:
# 预训练窗口大小不存在时的归一化。
# 如果预训练窗口大小不存在,使用当前窗口大小来归一化 relative_coords_table 中的坐标。
# 将高度坐标除以当前窗口高度减一。
relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
# 将宽度坐标除以当前窗口宽度减一。
relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
# 缩放坐标。将归一化的坐标缩放到范围[-8,8]。
relative_coords_table *= 8 # normalize to -8, 8
# torch.sign(input, *, out=None) -> Tensor
# torch.sign 是 PyTorch 中的一个函数,它返回输入张量中每个元素的符号。具体来说,这个函数会检查输入张量中的每个元素,并返回一个与输入张量形状相同的新张量,其中每个元素的值是: -1 如果原始元素是负数; 0 如果原始元素是零; 1 如果原始元素是正数。
# 参数说明 :
# input :输入的 Tensor ,可以是任意数值类型的张量。
# out :一个可选参数,用于输出结果的 Tensor 。如果提供,结果将被存储在这个 Tensor 中,而不是创建一个新的 Tensor 。
# 返回值 :
# 返回一个新的 Tensor ,其形状与输入 Tensor 相同,并且包含输入 Tensor 中每个元素的符号。
# torch.log2(input, *, out=None) -> Tensor
# torch.log2 是 PyTorch 中的一个函数,它计算输入张量中每个元素的以2为底的对数。这个函数对输入张量中的每个元素应用自然对数函数,然后将结果除以以2为底的自然对数(即 math.log(2) ),从而得到以2为底的对数。
# 参数说明 :
# input :输入的 Tensor ,必须是一个实数张量,因为对数函数只对正实数定义。
# out :一个可选参数,用于输出结果的 Tensor 。如果提供,结果将被存储在这个 Tensor 中,而不是创建一个新的 Tensor 。
# 返回值 :
# 返回一个新的 Tensor ,其形状与输入 Tensor 相同,并且包含输入 Tensor 中每个元素的以2为底的对数值。
# 计算对数尺度。
# torch.sign(relative_coords_table) :获取坐标的符号。 torch.abs(relative_coords_table) :获取坐标的绝对值。 torch.log2(...) :计算对数尺度。 np.log2(8) :计算对数的底数,这里为8。
relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
torch.abs(relative_coords_table) + 1.0) / np.log2(8)
# 注册缓冲区。将处理后的 relative_coords_table 注册为模块的缓冲区。这意味着它将被跟踪并用于模型的推理,但不会在训练中更新。
self.register_buffer("relative_coords_table", relative_coords_table)
# 获取窗口内每个标记的成对相对位置索引。
# get pair-wise relative position index for each token inside the window
# 生成坐标网格。
# 生成一个从0到 window_size[0] - 1 的张量,表示窗口内的高度坐标。
coords_h = torch.arange(self.window_size[0])
# 生成一个从0到 window_size[1] - 1 的张量,表示窗口内的宽度坐标。
coords_w = torch.arange(self.window_size[1])
# 使用 torch.meshgrid 生成坐标网格,然后使用 torch.stack 将高度和宽度坐标堆叠起来,形成一个形状为 (2, Wh, Ww) 的张量。
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww
# 展平坐标。将坐标网格展平,形成一个形状为 (2, Wh*Ww) 的张量。
coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww
# 计算相对坐标。
# 计算每个位置相对于其他位置的相对坐标,形成一个形状为 (2, Wh*Ww, Wh*Ww) 的张量。
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww
# 调整张量的维度顺序并确保它是连续的,形成一个形状为 (Wh*Ww, Wh*Ww, 2) 的张量。
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2
# 调整相对坐标。
# 将高度坐标调整为从0开始。
relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0
# 将宽度坐标调整为从0开始。
relative_coords[:, :, 1] += self.window_size[1] - 1
# 将高度坐标按窗口宽度的两倍减一进行缩放。
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
# 计算相对位置索引。
# 计算相对坐标的总和,得到一个形状为 (Wh*Ww, Wh*Ww) 的张量,表示每个位置相对于其他位置的相对位置索引。
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
# 注册缓冲区。将计算出的相对位置索引注册为模块的缓冲区,确保它在模型的整个生命周期中被正确管理。
self.register_buffer("relative_position_index", relative_position_index)
# 创建一个线性层,用于将输入特征的维度从 dim 转换为 dim * 3 。这个线性层将输入特征映射到查询(Q)、键(K)和值(V)三个向量,因此输出维度是输入维度的三倍。 bias=False 表示不使用偏置项。
self.qkv = nn.Linear(dim, dim * 3, bias=False)
if qkv_bias:
# 如果构造函数的参数 qkv_bias 为 True ,则为查询和值创建可学习的偏置参数。
# 创建一个形状为 (dim) 的全零张量,作为查询的偏置参数。
self.q_bias = nn.Parameter(torch.zeros(dim))
# 创建一个形状为 (dim) 的全零张量,作为值的偏置参数。
self.v_bias = nn.Parameter(torch.zeros(dim))
else:
# 如果 qkv_bias 为 False ,则不使用查询和值的偏置参数,将 self.q_bias 和 self.v_bias 设置为 None 。
self.q_bias = None
self.v_bias = None
# 创建一个dropout层,用于在计算注意力分数后应用dropout,以减少过拟合。 attn_drop 参数是dropout的概率。
self.attn_drop = nn.Dropout(attn_drop)
# 创建一个线性层,用于将自注意力机制的输出特征映射回原始维度 dim 。
self.proj = nn.Linear(dim, dim)
# 创建另一个dropout层,用于在自注意力机制的输出特征上应用dropout,以减少过拟合。 proj_drop 参数是dropout的概率。
self.proj_drop = nn.Dropout(proj_drop)
# 创建一个softmax层,用于将注意力分数转换为概率分布。 dim=-1 表示在最后一个维度上应用softmax,即在特征维度上。
self.softmax = nn.Softmax(dim=-1)
# self :指向类的实例。
# 1.x :输入数据。
# 2.mask :可选的 参数。
def forward(self, x, mask=None):
# 从输入数据 x 中获取其形状,并将其分解为三个变量: B_ (批次大小), N (序列长度),和 C (特征数量)。
B_, N, C = x.shape
# 初始化一个变量 qkv_bias ,用于存储查询(Q)、键(K)和值(V)的偏置项。
qkv_bias = None
# 检查是否存在查询的偏置项。
if self.q_bias is not None:
# 如果存在查询的偏置项,则创建一个新的偏置项张量 qkv_bias ,它将 查询的偏置项 、一个与值的偏置项形状相同但值为0的张量(不参与梯度计算),以及 值的偏置项 连接在一起。
qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
# torch.nn.functional.linear(input, weight, bias=None)
# torch.nn.functional.linear 是 PyTorch 中的一个函数,它提供了一个函数式接口来执行线性变换,这与 torch.nn.Linear 模块类似,但不需要创建一个线性层的实例。这个函数主要用于手动实现线性层,或者在没有定义具体模块的情况下使用。
# 参数说明 :
# input :输入的 Tensor ,形状为 (N, *, in_features) ,其中 N 是批次大小, * 表示任意数量的附加维度, in_features 是输入特征的数量。
# weight :权重矩阵,形状为 (out_features, in_features) ,其中 out_features 是输出特征的数量。
# bias :偏置向量,形状为 (out_features) 。这是一个可选参数,如果提供,将被添加到线性变换的结果中。
# 返回值 :
# 返回一个新的 Tensor ,形状为 (N, *, out_features) ,其中包含了输入数据经过线性变换后的结果。
# 使用线性变换(全连接层)对输入数据 x 进行变换,权重为 self.qkv.weight ,偏置项为 qkv_bias 。
qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
# 将线性变换后的张量 qkv 重塑并置换,以适应多头注意力机制的需要。 reshape 方法将张量的形状改变为 (B_, N, 3, self.num_heads, -1) ,其中3代表Q、K、V三个部分, permute 方法则重新排序这些维度。
qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
# 将置换后的张量 qkv 分解为三个部分:查询(Q)、键(K)和值(V)。这里使用索引来提取每个部分,并且这样做是为了使代码能够在TorchScript中运行,因为TorchScript不支持将张量作为元组使用。
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
# cosine attention 余弦注意力。
# 这里计算查询(Q)和键(K)的点积之前,先对它们进行归一化。 F.normalize 函数对输入张量沿着指定维度( dim=-1 ,即最后一个维度)进行L2范数归一化。
# 然后,通过矩阵乘法( @ 操作符)计算归一化后的Q和K的转置( .transpose(-2, -1) )的点积,得到注意力分数。
attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
# tensor.exp()
# 在 PyTorch 中, .exp() 是一个方法,它作用于 Tensor 对象,用于计算张量中每个元素的指数值。具体来说, .exp() 方法对张量中的每个元素应用指数函数 e^x ,其中 e 是自然对数的底数,大约等于 2.7182。
# 参数说明 :没有参数。
# 返回值 :
# 返回一个新的 Tensor ,其形状与原始 Tensor 相同,并且包含了原始 Tensor 中每个元素的指数值。
# self.logit_scale 是一个可学习的参数,用于缩放注意力分数以防止梯度消失或爆炸。
# torch.clamp 函数将 self.logit_scale 的值限制在一个范围内,这里的最大值是 log(1/0.01) ,即 -2 的对数值。然后,对限制后的值进行指数运算( .exp() ),得到最终的缩放因子 logit_scale 。
logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01))).exp()
# 将上一步得到的缩放因子应用到注意力分数上。
attn = attn * logit_scale
# self.cpb_mlp 是一个多层感知机(MLP),用于生成相对位置偏置。 self.relative_coords_table 包含了相对位置的坐标信息。MLP的输出被重新视图为 (-1, self.num_heads) ,即每个头一个偏置向量。
relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
# 根据 self.relative_position_index 索引 relative_position_bias_table ,得到每个位置的相对位置偏置,并将其重新视图以匹配窗口大小。
relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH
# 重新排列相对位置偏置的维度,并确保张量在内存中是连续的。
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww
# 对 相对位置偏置 应用sigmoid函数,然后乘以16以缩放偏置值。
relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
# 将相对位置偏置加到注意力分数上。 unsqueeze(0) 是在第一个维度(批次维度)上增加一个维度,以确保偏置张量的形状与注意力分数张量的形状匹配。
attn = attn + relative_position_bias.unsqueeze(0)
# 检查是否提供了掩码( mask )。掩码通常用于在注意力计算中屏蔽(忽略)某些元素,例如在处理序列时忽略填充(padding)元素。
if mask is not None:
# 如果提供了掩码,获取掩码的第一个维度的大小,这里假设掩码的形状是 (nW, N, N) ,其中 nW 是窗口的数量, N 是序列长度。
nW = mask.shape[0]
# 将注意力分数 attn 重塑为 (B_ // nW, nW, self.num_heads, N, N) 的形状,其中 B_ 是批次大小。
# 然后,将掩码 mask 在第1和第0维度上增加一个维度(即 unsqueeze(1) 和 unsqueeze(0) ),以便其形状与重塑后的 attn 相匹配,并将它们相加。
# 这样做的目的是将掩码的值(通常是非常大的负数)加到注意力分数上,使得这些位置在后续的softmax归一化中接近于0。
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
# 将经过掩码处理的注意力分数 attn 重塑回 (-1, self.num_heads, N, N) 的形状,以便于进行softmax归一化。
attn = attn.view(-1, self.num_heads, N, N)
# 应用softmax函数对注意力分数进行归一化,使得每个头的每个位置的注意力分数之和为1。
attn = self.softmax(attn)
else:
# 如果没有提供掩码,直接对原始的注意力分数应用softmax归一化。
attn = self.softmax(attn)
# 在softmax归一化之后,应用注意力dropout。这是一个正则化技术,用于随机将一些注意力分数设置为0,以防止过拟合。
attn = self.attn_drop(attn)
# 开始一个 try 块,用于尝试执行可能抛出异常的代码。
try:
# 使用矩阵乘法( @ 操作符)将注意力权重( attn )与值( v )相乘,得到加权的值。
# .transpose(1, 2) 将结果张量的第一个和第二个维度交换,通常是将头的维度( self.num_heads )和序列长度( N )交换。
# .reshape(B_, N, C) 将结果张量重塑为原始输入的形状,其中 B_ 是批次大小, N 是序列长度, C 是特征数量。
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
except:
# 如果在 try 块中的操作抛出异常,将执行 except 块中的代码。
# 如果原始操作失败,尝试将 attn 和 v 都转换为半精度浮点数( half ),然后执行相同的矩阵乘法、转置和重塑操作。这通常用于处理内存不足或计算资源受限的情况,因为半精度浮点数占用的内存更少。
x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
# 将加权的值通过一个线性层( self.proj ),这个线性层将特征维度从 C 映射回原始的特征维度。
x = self.proj(x)
# 对线性投影后的结果应用dropout( self.proj_drop ),这是一种正则化技术,用于减少过拟合。
x = self.proj_drop(x)
# 返回最终的结果 x ,它是经过自注意力机制处理后的特征表示。
return x
# 这段代码定义了一个名为 extra_repr 的方法,它是 Python 类中的特殊方法,用于提供类的额外字符串表示。这个方法通常被用来返回类的配置或状态信息,以便在打印类实例时提供更多的上下文信息。
def extra_repr(self) -> str:
# 在这个特定的 extra_repr 方法中,返回的字符串包含了以下属性的值 :
# dim :表示输入特征的维度。
# window_size :表示窗口的尺寸。
# pretrained_window_size :表示预训练模型中使用的窗口尺寸。
# num_heads :表示多头注意力机制中的头数。
return f'dim={self.dim}, window_size={self.window_size}, ' \
f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}'
# 这段代码定义了一个名为 flops 的方法,用于计算在给定序列长度 N 的情况下,某个操作或模块(很可能是一个自注意力模块)的浮点运算次数(FLOPs,即每秒浮点运算次数)。FLOPs 是衡量算法或模型计算复杂度的一个指标,通常用于评估算法的效率。
# self :指向类的实例。
# 1.N :序列长度。
def flops(self, N):
# 计算 1 个窗口的浮点运算次数,标记长度为 N。
# calculate flops for 1 window with token length of N
# 初始化一个变量 flops 来累计总的浮点运算次数。
flops = 0
# qkv = self.qkv(x)
# 计算 QKV(Query, Key, Value)矩阵乘法的浮点运算次数。这里假设 self.qkv 是一个线性层,它将输入 x 映射到三个矩阵(Q、K、V),每个矩阵的维度都是 N x self.dim 。因此,每个矩阵的乘法运算次数是 N * self.dim * self.dim ,由于有三个矩阵,所以乘以3。
flops += N * self.dim * 3 * self.dim
# attn = (q @ k.transpose(-2, -1))
# 计算注意力分数的浮点运算次数。这里假设 Q 和 K 的维度分别是 N x (self.dim // self.num_heads) ,并且它们通过矩阵乘法计算注意力分数,所以运算次数是 N * (self.dim // self.num_heads) * N 。由于有 self.num_heads 个头,所以乘以 self.num_heads 。
flops += self.num_heads * N * (self.dim // self.num_heads) * N
# x = (attn @ v)
# 计算将注意力分数与 V 相乘的浮点运算次数。这里假设 V 的维度是 N x (self.dim // self.num_heads) ,所以运算次数是 N * N * (self.dim // self.num_heads) 。同样,由于有 self.num_heads 个头,所以乘以 self.num_heads 。
flops += self.num_heads * N * N * (self.dim // self.num_heads)
# x = self.proj(x)
# 计算最后一个线性层(投影层)的浮点运算次数。这里假设输入的维度是 N x self.dim ,输出的维度也是 N x self.dim ,所以运算次数是 N * self.dim * self.dim 。
flops += N * self.dim * self.dim
# 返回计算出的总浮点运算次数。
return flops
# 这段代码定义了一个名为 Mlp_v2 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了一个两层的多层感知机(MLP),通常用于深度学习中的各种任务,如特征变换或分类。
class Mlp_v2(nn.Module):
# 1.in_features :输入特征的数量。
# 2.hidden_features :第一个线性层(fc1)的输出特征数量。如果未指定,则默认与输入特征数量相同。
# 3.out_features :第二个线性层(fc2)的输出特征数量。如果未指定,则默认与输入特征数量相同。
# 4.act_layer :激活函数,默认为 nn.SiLU (也称为 Swish)。
# 5.drop :Dropout 比率,默认为 0,表示不应用 Dropout。
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.):
super().__init__()
# 如果 out_features 未指定,则设置为 in_features 。
out_features = out_features or in_features
# 如果 hidden_features 未指定,则设置为 in_features 。
hidden_features = hidden_features or in_features
# 定义第一个线性层,将输入特征从 in_features 映射到 hidden_features 。
self.fc1 = nn.Linear(in_features, hidden_features)
# 实例化激活函数。
self.act = act_layer()
# 定义第二个线性层,将特征从 hidden_features 映射到 out_features 。
self.fc2 = nn.Linear(hidden_features, out_features)
# 定义 Dropout 层,Dropout 比率为 drop 。
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
# 这段代码定义了一个名为 window_partition_v2 的函数,它用于将输入张量 x 分割成多个窗口。这种操作在视觉任务中,特别是在 Vision Transformer(ViT)模型中,用于将图像划分为非重叠的局部区域(窗口),以便在每个窗口上独立地应用自注意力机制。
# 1.x :输入张量,形状为 (B, H, W, C) ,其中 B 是批次大小, H 和 W 分别是张量的高度和宽度, C 是通道数。
# 2.window_size :窗口的大小,是一个整数,表示每个窗口的高度和宽度。
def window_partition_v2(x, window_size):
# 从输入张量 x 中获取其形状,并将其分解为 B (批次大小)、 H (高度)、 W (宽度)和 C (通道数)。
B, H, W, C = x.shape
# 使用 view 方法将输入张量 x 重塑为 (B, H // window_size, window_size, W // window_size, window_size, C) 的形状。这一步实际上是将高度和宽度分别按照 window_size 进行分割,形成多个 window_size x window_size 的小块。
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
# permute(0, 1, 3, 2, 4, 5) :重新排列张量的维度,将 (B, H // window_size, window_size, W // window_size, window_size, C) 变为 (B, H // window_size, W // window_size, window_size, window_size, C) 。这样做是为了将窗口的行和列分开处理。
# contiguous() :确保张量在内存中是连续的,这是在进行 view 操作之前的必要步骤,以避免维度不匹配的问题。
# view(-1, window_size, window_size, C) :将重新排列后的张量重塑为 (-1, window_size, window_size, C) 的形状,其中 -1 表示自动计算该维度的大小,以确保总元素数量不变。
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
# 返回分割后的窗口张量,形状为 (-1, window_size, window_size, C) ,其中 -1 表示总共有 B * (H // window_size) * (W // window_size) 个窗口。
return windows
# 这个函数是实现 Vision Transformer 模型中的一个重要步骤,它将输入张量分割成多个窗口,以便在每个窗口上独立地应用自注意力机制。这种方法可以减少计算量,同时保持模型对局部特征的捕捉能力。
# 这段代码定义了一个名为 window_reverse_v2 的函数,它用于将分割后的窗口张量重新组合回原始的张量形状。这通常是在自注意力机制处理完每个窗口后,将窗口特征合并回整个特征图的操作。
# 1.windows :经过 window_partition_v2 函数分割后的窗口张量,形状为 (-1, window_size, window_size, C) ,其中 -1 表示窗口的总数。
# 2.window_size :窗口的大小,与 window_partition_v2 函数中的相同。
# 3.H 和 4.W :原始张量的高度和宽度。
def window_reverse_v2(windows, window_size, H, W):
# 计算批次大小 B 。这里假设 windows 的第一个维度是所有窗口的总数,通过将窗口总数除以原始张量的每个维度上的窗口数( H // window_size 和 W // window_size ),可以得到批次大小。
B = int(windows.shape[0] / (H * W / window_size / window_size))
# 使用 view 方法将窗口张量 windows 重塑为 (B, H // window_size, W // window_size, window_size, window_size, -1) 的形状。这一步是为了将窗口重新排列成原始张量的高度和宽度的网格形式。
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
# permute(0, 1, 3, 2, 4, 5) :重新排列张量的维度,将 (B, H // window_size, W // window_size, window_size, window_size, -1) 变为 (B, H // window_size, window_size, W // window_size, window_size, -1) 。这样做是为了将窗口的行和列重新组合成原始的高度和宽度。
# contiguous() :确保张量在内存中是连续的,这是在进行 view 操作之前的必要步骤。
# view(B, H, W, -1) :将重新排列后的张量重塑为 (B, H, W, -1) 的形状,其中 -1 表示自动计算该维度的大小,以确保总元素数量不变。
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
# 返回重新组合后的张量,形状为 (B, H, W, C) ,其中 C 是通道数,与原始张量的形状相同。
return x
# 这个函数是 window_partition_v2 函数的逆操作,它将分割后的窗口特征合并回整个特征图,以便进行后续的处理或输出。这种分割和合并操作在 Vision Transformer 和其他需要处理局部特征的模型中非常常见。
# 这段代码定义了一个名为 SwinTransformerLayer_v2 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了 Swin Transformer 的一个单独的层,Swin Transformer 是一种用于计算机视觉任务的 Transformer 架构。
class SwinTransformerLayer_v2(nn.Module):
# 1.dim :输入特征的维度。
# 2.num_heads :多头自注意力机制中的头数。
# 3.window_size :用于自注意力计算的窗口大小,默认为 7。
# 4.shift_size :在 Swin Transformer 中,用于计算窗口划分时的偏移量,默认为 0。
# 5.mlp_ratio :多层感知机(MLP)的隐藏层大小与输入特征维度的比例,默认为 4。
# 6.qkv_bias :是否在 QKV(Query, Key, Value)线性层中使用偏置项,默认为 True。
# 7.drop :输入特征的 Dropout 比率,默认为 0。
# 8.attn_drop :注意力分数的 Dropout 比率,默认为 0。•
# 9.drop_path :随机丢弃路径(Stochastic Depth)的比率,默认为 0。
# 10.act_layer :激活函数,默认为 nn.SiLU (也称为 Swish)。
# 11.norm_layer :归一化层,默认为 nn.LayerNorm 。
# 12.pretrained_window_size :预训练时使用的窗口大小,默认为 0,表示与 window_size 相同。
def __init__(self, dim, num_heads, window_size=7, shift_size=0,
mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,
act_layer=nn.SiLU, norm_layer=nn.LayerNorm, pretrained_window_size=0):
# 调用父类 nn.Module 的构造函数,这是初始化 PyTorch 模块的标准做法。
super().__init__()
# 输入特征的维度。这个属性被设置为构造函数参数 dim 的值,它表示每个输入特征向量的维数。
self.dim = dim
#self.input_resolution = input_resolution
# 多头自注意力机制中的头数。这个属性被设置为构造函数参数 num_heads 的值,它决定了自注意力计算中将使用多少个注意力头。
self.num_heads = num_heads
# 用于自注意力计算的窗口大小。这个属性被设置为构造函数参数 window_size 的值,它定义了在自注意力计算中每个窗口的高度和宽度。
self.window_size = window_size
# 在 Swin Transformer 中,用于计算窗口划分时的偏移量。这个属性被设置为构造函数参数 shift_size 的值,它用于实现窗口划分时的位移操作,以便在不同的窗口之间共享信息。
self.shift_size = shift_size
# 多层感知机(MLP)的隐藏层大小与输入特征维度的比例。这个属性被设置为构造函数参数 mlp_ratio 的值,它定义了 MLP 中隐藏层的大小相对于输入特征维度的缩放比例。
self.mlp_ratio = mlp_ratio
#if min(self.input_resolution) <= self.window_size:
# # if window size is larger than input resolution, we don't partition windows
# self.shift_size = 0
# self.window_size = min(self.input_resolution)
# 这是一个断言语句,用于确保 shift_size 的值在合理的范围内,即 0 到 window_size (不包括 window_size )。如果 shift_size 的值不满足这个条件,将会抛出一个 AssertionError 。
assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
# 归一化层。
# 初始化第一个归一化层 norm1 ,通常用于自注意力之前的归一化。
self.norm1 = norm_layer(dim)
# 自注意力机制。
# 初始化一个 WindowAttention_v2 模块,这是一个窗口化的自注意力机制,它在局部窗口内计算自注意力。这里传入了多个参数,包括 :
# 特征维度 dim 、窗口大小 (self.window_size, self.window_size) 、头数 num_heads 、是否使用偏置 qkv_bias 、注意力 Dropout 比率 attn_drop 、投影 Dropout 比率 proj_drop ,以及预训练时的窗口大小 (pretrained_window_size, pretrained_window_size) 。
self.attn = WindowAttention_v2(
dim, window_size=(self.window_size, self.window_size), num_heads=num_heads,
qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,
pretrained_window_size=(pretrained_window_size, pretrained_window_size))
# 随机丢弃路径。
# 如果 drop_path 大于 0 ,则初始化一个 DropPath 模块,这是一种正则化技术,用于在训练过程中随机丢弃整个路径。如果 drop_path 不大于 0 ,则使用 nn.Identity ,即一个恒等映射,这意味着不会应用任何丢弃。
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
# 初始化第二个归一化层 norm2 ,通常用于 MLP 之后的归一化。
self.norm2 = norm_layer(dim)
# 多层感知机(MLP)。计算 MLP 的隐藏层维度,这是基于输入特征维度 dim 和比例 mlp_ratio 计算得出的。
mlp_hidden_dim = int(dim * mlp_ratio)
# 初始化一个 Mlp_v2 模块,这是一个两层的 MLP,它包含一个激活函数和一个 Dropout 层。这里传入了输入特征维度 dim 、隐藏层维度 mlp_hidden_dim 、激活函数 act_layer 和 Dropout 比率 drop 。
self.mlp = Mlp_v2(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
# 这段代码定义了一个名为 create_mask 的方法,它用于创建 Swin Transformer 中的自注意力掩码(attention mask)。这个掩码用于控制自注意力机制中的注意力分布,特别是在实现 Swin Transformer 的 shifted window multi-head self-attention (SW-MSA) 时。
# 1.H 和 2.W :输入特征图的高度和宽度。
def create_mask(self, H, W):
# calculate attention mask for SW-MSA
# 创建一个形状为 (1, H, W, 1) 的全零张量,用于存储掩码值。
img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1
# self.window_size 和 self.shift_size 是模型中的参数,分别代表 窗口的大小 和 窗口之间的偏移量 。
# h_slices 和 w_slices : 这两个元组分别用于处理高度和宽度方向上的切片。
# 每个元组包含三个切片对象,对应于不同的窗口区域 :
# 第一个切片 ( slice(0, -self.window_size) ) :表示从当前窗口的开始到窗口大小之前的位置。
# 第二个切片 ( slice(-self.window_size, -self.shift_size) ) :表示从窗口大小之前的位置到位移大小之前的位置。
# 第三个切片 ( slice(-self.shift_size, None) ) :表示从位移大小之前的位置到窗口的结束。
h_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
w_slices = (slice(0, -self.window_size),
slice(-self.window_size, -self.shift_size),
slice(-self.shift_size, None))
# 初始化一个计数器,用于为每个窗口分配一个唯一的标识符。
cnt = 0
# 循环遍历每个切片。
# 遍历每个水平和垂直切片的组合。
for h in h_slices:
for w in w_slices:
# 在 img_mask 中设置当前切片的值为计数器的值。
img_mask[:, h, w, :] = cnt
# 递增计数器。
cnt += 1
# def window_partition(x, window_size): -> 它用于将输入的四维特征图 x 分割成多个窗口。返回 分割后的窗口 张量。 -> return windows
# 使用 window_partition 函数将 img_mask 分割成窗口大小的小块。
mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1
# 将分割后的窗口张量重塑,以便于后续操作。
mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
# 这行代码计算了自注意力掩码。
# 首先, unsqueeze 操作在指定维度上增加一个维度。 unsqueeze(1) 和 unsqueeze(2) 分别在第二个和第三个维度上增加维度,使得张量的形状变为 (-1, 1, self.window_size * self.window_size) 和 (-1, self.window_size * self.window_size, 1) 。
# 然后,通过相减操作,对于每个窗口内的元素,如果它们属于同一个窗口(即它们的 mask_windows 值相同),则相减结果为0;如果它们属于不同窗口,则相减结果非0。这样,我们可以得到一个表示元素间是否属于同一窗口的掩码。
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
# masked_fill(mask, value)
# masked_fill_(mask, value)
# masked_fill 是 PyTorch 中的一个函数,它用于根据给定的掩码(mask)将张量(tensor)中指定位置的元素替换为特定的值。
# 参数 :
# mask :一个布尔类型的张量(BoolTensor),其形状必须与要操作的张量形状相同或者可以广播到张量的形状。在 mask 中为 True 的位置将被替换,为 False 的位置保持不变。
# value :一个标量或与被操作的张量具有相同形状的张量,用于指定替换的值。
# 功能描述 :
# masked_fill 函数会检查 mask 中的每个元素,如果 mask 对应的位置为 True ,则在输出张量中用 value 替换原张量中对应位置的值;如果为 False ,则保留原张量中对应位置的值。
# 返回值 :
# 返回一个新的张量,其中包含了在 mask 指定位置上被替换的值,其余位置保持原样。原始张量不会被改变( masked_fill 是非原地操作)。
# 使用 masked_fill 函数填充掩码。非零值被替换为一个非常大的负数(通常用于 softmax 之前的掩码),零值被替换为零。
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
# 返回计算出的自注意力掩码。
return attn_mask
# 这个掩码可以用于自注意力机制中,通过将掩码乘以注意力分数,使得某些位置的注意力权重接近于零,从而实现对特定位置的忽略。这种方法在处理图像或其他序列数据时特别有用,因为它允许模型在不同的窗口和位置之间共享信息。
# x :输入张量,形状为 (B, C, H, W) ,其中 B 是批次大小, C 是通道数, H 是高度, W 是宽度。
def forward(self, x):
# reshape x[b c h w] to x[b l c]
# 从输入张量 x 中获取其形状,并将其分解为 H_ (高度)和 W_ (宽度)。
_, _, H_, W_ = x.shape
# 初始化一个标志变量 Padding ,用于指示是否需要对输入张量进行填充。
Padding = False
# 判断是否需要填充。如果输入张量的高度或宽度小于窗口大小,或者高度和宽度不能被窗口大小整除,则需要填充。
if min(H_, W_) < self.window_size or H_ % self.window_size!=0 or W_ % self.window_size!=0:
Padding = True
# print(f'img_size {min(H_, W_)} is less than (or not divided by) window_size {self.window_size}, Padding.')
# 计算填充量。
# 计算宽度方向需要填充的像素数。
pad_r = (self.window_size - W_ % self.window_size) % self.window_size
# 计算高度方向需要填充的像素数。
pad_b = (self.window_size - H_ % self.window_size) % self.window_size
# 使用 PyTorch 的 F.pad 函数对输入张量 x 进行填充。填充是在宽度的右侧和高度的底部进行的。
x = F.pad(x, (0, pad_r, 0, pad_b))
# print('2', x.shape)
# 输入张量重塑。获取输入张量 x 的形状,分别代表批次大小(B)、通道数(C)、高度(H)和宽度(W)。
B, C, H, W = x.shape
# 计算张量的序列长度 L ,即张量的面积(高度乘以宽度)。
L = H * W
# 将输入张量 x 的维度从 (B, C, H, W) 置换为 (B, H, W, C) 。 调用 contiguous 确保张量在内存中是连续的。 将张量重塑为 (B, L, C) ,其中 L 是序列长度, C 是通道数,这样每个序列元素对应于输入特征图的一个元素。
x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C) # b, L, c
# create mask from init to forward
# 自注意力掩码创建。
# 如果 shift_size 大于 0,表示需要进行窗口位移操作,因此需要创建一个掩码。
if self.shift_size > 0:
# def create_mask(self, H, W): -> 它用于创建 Swin Transformer 中的自注意力掩码(attention mask)。返回计算出的自注意力掩码。 -> return attn_mask
# 调用 create_mask 方法创建掩码,并确保掩码在与输入张量相同的设备上(例如 GPU 或 CPU)。
attn_mask = self.create_mask(H, W).to(x.device)
else:
# 如果 shift_size 不大于 0,不需要位移,因此掩码为 None 。
attn_mask = None
# 快捷连接。创建一个快捷连接 shortcut ,它将用于将输入 x 加到自注意力层的输出上,以便进行残差连接。
shortcut = x
# 输入张量重塑为原始形状。将张量 x 从 (B, L, C) 重塑回 (B, H, W, C) ,以便进行循环位移操作。
x = x.view(B, H, W, C)
# cyclic shift 循环位移。
# 如果 shift_size 大于 0,执行循环位移操作。
if self.shift_size > 0:
# torch.roll(input, shifts, dims=None)
# torch.roll 是 PyTorch 中的一个函数,它对输入张量进行滚动操作,也就是将张量中的元素沿着指定的维度“滚动”一定数量的位置。这个操作会将指定维度上的元素按照指定的步数向前或向后移动,移动出的元素会被重新放置到另一端。
# 参数说明 :
# input :输入的 Tensor ,即要进行滚动操作的张量。
# shifts :一个整数或整数元组,指定每个维度上元素需要滚动的步数。如果是一个整数,则在所有维度上滚动相同的步数。
# dims :一个整数或整数元组,指定要滚动的维度。如果为 None ,则滚动所有维度。
# 返回值 :
# 返回一个新的 Tensor ,其形状与输入 Tensor 相同,但元素已经根据指定的步数和维度进行了滚动。
# 使用 torch.roll 函数对张量 x 进行循环位移,位移量为 (self.shift_size, self.shift_size) ,分别在高度和宽度维度上。
shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
else:
# 如果 shift_size 不大于 0,不进行位移, shifted_x 就是原始的 x 。
shifted_x = x
# partition windows
# 分割窗口。
# def window_partition_v2(x, window_size):
# -> 它用于将输入张量 x 分割成多个窗口。返回分割后的窗口张量,形状为 (-1, window_size, window_size, C) ,其中 -1 表示总共有 B * (H // window_size) * (W // window_size) 个窗口。
# -> return windows
# 使用 window_partition_v2 函数将循环位移后的张量 shifted_x 分割成多个窗口。每个窗口的大小为 self.window_size x self.window_size ,结果张量的形状为 (nW*B, window_size, window_size, C) ,其中 nW 是窗口的总数, B 是批次大小, C 是通道数。
x_windows = window_partition_v2(shifted_x, self.window_size) # nW*B, window_size, window_size, C
# 将分割后的窗口张量 x_windows 重塑,以便于输入到自注意力模块。重塑后的形状为 (nW*B, window_size*window_size, C) 。
x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C
# W-MSA/SW-MSA
# 应用自注意力机制。
# 调用 self.attn (自注意力模块)处理重塑后的窗口张量 x_windows ,并传入之前创建的掩码 attn_mask 。输出的形状为 (nW*B, window_size*window_size, C) 。
attn_windows = self.attn(x_windows, mask=attn_mask) # nW*B, window_size*window_size, C
# merge windows
# 重新组合窗口。将自注意力模块的输出 attn_windows 重塑回窗口的形式,形状为 (nW*B, window_size, window_size, C) 。
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
# def window_reverse_v2(windows, window_size, H, W):
# -> 它用于将分割后的窗口张量重新组合回原始的张量形状。这通常是在自注意力机制处理完每个窗口后,将窗口特征合并回整个特征图的操作。返回重新组合后的张量,形状为 (B, H, W, C) ,其中 C 是通道数,与原始张量的形状相同。
# -> return x
# 使用 window_reverse_v2 函数将自注意力处理后的窗口张量 attn_windows 重新组合回原始的张量形状 (B, H', W', C) ,其中 H' 和 W' 分别是可能经过填充的高度和宽度。
shifted_x = window_reverse_v2(attn_windows, self.window_size, H, W) # B H' W' C
# 逆循环位移。
# reverse cyclic shift
# 如果 shift_size 大于 0,表示之前进行了循环位移,现在需要进行逆操作。
if self.shift_size > 0:
# 使用 torch.roll 函数将 shifted_x 沿着高度和宽度维度向后滚动 self.shift_size 步,以抵消之前的位移。
x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
else:
# 如果 shift_size 不大于 0,不需要进行位移,直接使用 shifted_x 。
x = shifted_x
# 残差连接和归一化。
# 将张量 x 重塑为 (B, H * W, C) ,以便于进行残差连接。
x = x.view(B, H * W, C)
# 将归一化和可能的dropout后的 x 与快捷连接 shortcut 相加,形成残差连接。
x = shortcut + self.drop_path(self.norm1(x))
# FFN
# 多层感知机(MLP)。
# 应用第二个归一化层 norm2 和 MLP mlp 到 x ,然后通过 drop_path 进行可能的dropout,最后与原始 x 相加,形成另一个残差连接。
x = x + self.drop_path(self.norm2(self.mlp(x)))
# 输出重塑。将张量 x 的维度置换为 (0, 2, 1) ,即 (B, C, H * W) 。 调用 contiguous 确保张量在内存中是连续的。 将张量重塑为 (-1, C, H, W) ,恢复到原始的四维形状。
x = x.permute(0, 2, 1).contiguous().view(-1, C, H, W) # b c h w
# 填充撤销。将张量 x 裁剪回原始的高度 H_ 和宽度 W_ ,撤销之前为了适应窗口大小而进行的填充。
if Padding:
x = x[:, :, :H_, :W_] # reverse padding 反向填充。
# 返回最终的输出张量 x 。
return x
# 这段代码定义了一个名为 extra_repr 的方法,它是 Python 类中的特殊方法,用于提供类的额外字符串表示。这个方法通常被用来返回类的配置或状态信息,以便在打印类实例时提供更多的上下文信息。
def extra_repr(self) -> str:
# 在这个特定的 extra_repr 方法中,返回的字符串包含了以下属性的值 :
# dim :类的某个属性,表示特征的维度。
# input_resolution :类的某个属性,表示输入数据的分辨率或尺寸。
# num_heads :类的某个属性,表示多头自注意力机制中的头数。
# window_size :类的某个属性,表示自注意力计算中使用的窗口大小。
# shift_size :类的某个属性,表示在计算自注意力时窗口的位移大小。
# mlp_ratio :类的某个属性,表示多层感知机(MLP)中隐藏层大小与输入特征维度的比例。
return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
# 这段代码定义了一个名为 flops 的方法,用于计算 Swin Transformer 层中的浮点运算次数(FLOPs)。FLOPs 是衡量算法或模型计算复杂度的一个指标,通常用于评估算法的效率。
def flops(self):
# 初始化一个变量 flops 来累计总的浮点运算次数。
flops = 0
# 获取类的 input_resolution 属性,它包含了输入特征图的高度 H 和宽度 W 。
H, W = self.input_resolution
# norm1 注释表明接下来的计算是归一化层 norm1 的 FLOPs。
# 计算归一化层 norm1 的 FLOPs。每个元素都需要进行一次归一化操作,因此 FLOPs 为输入特征维度 dim 乘以特征图的面积 H * W 。
flops += self.dim * H * W
# W-MSA/SW-MSA 注释表明接下来的计算是窗口自注意力机制(W-MSA 或 SW-MSA)的 FLOPs。
# 计算窗口的总数 nW ,即输入特征图的面积除以窗口大小的平方。
nW = H * W / self.window_size / self.window_size
# def flops(self, N): -> 用于计算在给定序列长度 N 的情况下,某个操作或模块(很可能是一个自注意力模块)的浮点运算次数(FLOPs,即每秒浮点运算次数)。返回计算出的总浮点运算次数。 -> return flops
# 计算每个窗口内自注意力机制的 FLOPs,并乘以窗口的总数。这里假设 self.attn 有一个 flops 方法可以计算单个窗口的 FLOPs。
flops += nW * self.attn.flops(self.window_size * self.window_size)
# mlp 注释表明接下来的计算是多层感知机(MLP)的 FLOPs。
# 计算 MLP 的 FLOPs。MLP 包含两个线性层,每个线性层的 FLOPs 为输入特征维度 dim 乘以输出特征维度 dim (对于第一个线性层)和 dim 乘以隐藏层维度(对于第二个线性层)。
# 由于隐藏层维度是输入维度的 mlp_ratio 倍,所以总 FLOPs 为 2 * H * W * self.dim * self.dim * self.mlp_ratio 。
flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
# norm2 注释表明接下来的计算是归一化层 norm2 的 FLOPs。
# 计算归一化层 norm2 的 FLOPs,与 norm1 类似。
flops += self.dim * H * W
# 返回计算出的总 FLOPs。
return flops
# 这个方法提供了一个估计 Swin Transformer 层计算复杂度的方式,有助于在设计和优化模型时考虑计算资源的需求。通过计算 FLOPs,我们可以了解模型在不同输入大小下的计算资源需求,这对于优化模型性能和选择适合特定硬件的模型架构非常有用。
# 这段代码定义了一个名为 SwinTransformer2Block 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了 Swin Transformer 的一个变体,其中包含了一个可选的卷积层和多个 Swin Transformer 层。
class SwinTransformer2Block(nn.Module):
# c1 :输入特征的通道数。
# c2 :输出特征的通道数。
# num_heads :多头自注意力机制中的头数。
# num_layers :Swin Transformer 层的数量。
# window_size :用于自注意力计算的窗口大小,默认为 7。
def __init__(self, c1, c2, num_heads, num_layers, window_size=7):
super().__init__()
# 初始化一个标志变量 self.conv ,用于指示是否需要卷积层。
self.conv = None
# 如果输入和输出通道数不同,则需要一个卷积层来调整通道数。
if c1 != c2:
# 创建一个卷积层,将输入特征从 c1 通道映射到 c2 通道。
self.conv = Conv(c1, c2)
# remove input_resolution 删除 input_resolution 。
# 创建一个序列容器 self.blocks ,包含 num_layers 个 SwinTransformerLayer_v2 层。每个层的 shift_size 根据其索引 i 决定,如果 i 是偶数,则 shift_size 为 0;如果 i 是奇数,则 shift_size 为 window_size // 2 。
self.blocks = nn.Sequential(*[SwinTransformerLayer_v2(dim=c2, num_heads=num_heads, window_size=window_size,
shift_size=0 if (i % 2 == 0) else window_size // 2) for i in range(num_layers)])
def forward(self, x):
# 如果存在卷积层,则对输入 x 应用该卷积层。
if self.conv is not None:
x = self.conv(x)
# 将输入 x 通过 self.blocks 序列中的所有 Swin Transformer 层。
x = self.blocks(x)
return x
# 这个类实现了 Swin Transformer 的一个块,其中可以包含一个卷积层来调整通道数,以及多个 Swin Transformer 层来处理特征。
# 这种设计允许模型在不同的阶段调整特征的通道数,并在不同的层之间交替使用不同的自注意力模式(通过改变 shift_size )。这种交替模式有助于模型捕获不同尺度的特征,并提高模型的表示能力。
# 这段代码定义了一个名为 ST2CSPA 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了一个结合了 Cross Stage Partial Networks (CSP) 和 Swin Transformer 的网络结构。
class ST2CSPA(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
# 1.c1 :输入特征的通道数。
# 2.c2 :输出特征的通道数。
# 3.n :Swin Transformer 块的数量,默认为 1。
# 4.shortcut :是否使用快捷连接,默认为 True。
# 5.g :分组卷积的组数,默认为 1。
# 6.e :扩张系数,用于计算隐藏层通道数,默认为 0.5。
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(ST2CSPA, self).__init__()
# 计算隐藏层的通道数 c_ ,它是输出通道数 c2 的 e 倍。
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c1, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
# 计算多头自注意力机制中的头数,这里简单地将 c_ 除以 32。
num_heads = c_ // 32
# 创建一个 SwinTransformer2Block 实例 m ,它包含了 n 个 Swin Transformer 层。
self.m = SwinTransformer2Block(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
y1 = self.m(self.cv1(x))
y2 = self.cv2(x)
# 将 y1 和 y2 在通道维度上合并,然后通过 cv3 卷积层输出最终结果。
return self.cv3(torch.cat((y1, y2), dim=1))
# 这个类结合了 CSP 网络的设计理念和 Swin Transformer 的自注意力机制,旨在提高网络的特征提取能力和效率。通过这种方式,模型可以有效地处理图像特征,并在不同的阶段调整特征的表示。
# 这段代码定义了一个名为 ST2CSPB 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了一个结合了 Cross Stage Partial Networks (CSP) 和 Swin Transformer 的网络结构,具体来说是 CSP 瓶颈(Bottleneck)变体。
class ST2CSPB(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(ST2CSPB, self).__init__()
# 计算隐藏层的通道数 c_ ,这里直接使用输出通道数 c2 。
c_ = int(c2) # hidden channels
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_, c_, 1, 1)
self.cv3 = Conv(2 * c_, c2, 1, 1)
num_heads = c_ // 32
self.m = SwinTransformer2Block(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
x1 = self.cv1(x)
y1 = self.m(x1)
y2 = self.cv2(x1)
return self.cv3(torch.cat((y1, y2), dim=1))
# 这个类结合了 CSP 网络的设计理念和 Swin Transformer 的自注意力机制,旨在提高网络的特征提取能力和效率。通过这种方式,模型可以有效地处理图像特征,并在不同的阶段调整特征的表示。
# 与 ST2CSPA 类似, ST2CSPB 也是 CSP 结构的一个变体,但具体的实现细节有所不同。
# 这段代码定义了一个名为 ST2CSPC 的类,它是一个神经网络模块,继承自 PyTorch 的 nn.Module 。这个类实现了一个结合了 Cross Stage Partial Networks (CSP) 和 Swin Transformer 的网络结构,具体来说是 CSP 瓶颈(Bottleneck)变体 C。
class ST2CSPC(nn.Module):
# CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
super(ST2CSPC, self).__init__()
c_ = int(c2 * e) # hidden channels
# 创建一个卷积层 cv1 ,用于将输入特征从 c1 通道映射到 c_ 通道,卷积核大小为 1x1。
self.cv1 = Conv(c1, c_, 1, 1)
# 创建另一个卷积层 cv2 ,用于在 CSP 结构中进一步处理特征,卷积核大小同样为 1x1。
self.cv2 = Conv(c1, c_, 1, 1)
# 创建一个卷积层 cv3 ,用于在 CSP 结构中进一步处理特征,卷积核大小为 1x1。
self.cv3 = Conv(c_, c_, 1, 1)
# 创建一个卷积层 cv4 ,用于将两个 c_ 通道的特征合并,并映射到 c2 通道。
self.cv4 = Conv(2 * c_, c2, 1, 1)
num_heads = c_ // 32
self.m = SwinTransformer2Block(c_, c_, num_heads, n)
#self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
def forward(self, x):
# 将输入 x 通过 cv1 卷积层,然后通过 Swin Transformer 块 m ,最后通过 cv3 卷积层。
y1 = self.cv3(self.m(self.cv1(x)))
y2 = self.cv2(x)
return self.cv4(torch.cat((y1, y2), dim=1))
# 这个类结合了 CSP 网络的设计理念和 Swin Transformer 的自注意力机制,旨在提高网络的特征提取能力和效率。通过这种方式,模型可以有效地处理图像特征,并在不同的阶段调整特征的表示。
# ST2CSPC 是 CSP 结构的一个变体,其中包含了一个额外的卷积层 cv3 ,用于处理经过 Swin Transformer 块的特征。
##### end of swin transformer v2 #####