计算机视觉方面的一些模块
# __all__ 是一个可选的列表，定义在模块级别。当使用 from ... import * 语句时，如果模块中定义了
# __all__，则只有 __all__ 列表中的名称会被导入。这是模块作者控制哪些公开API被导入的一种方式。
# 使用 * 导入的行为
# 如果模块中有 __all__ 列表：只有 __all__ 列表中的名称会被导入
# 如果模块中没有 __all__ 列表：Python 解释器会尝试导入模块中定义的所有公共名称（即不
# 是以下划线 _ 开头的名称）。但是，这通常不包括以单下划线或双下划线开头的特殊方法或变量
def autopad(k, p=None, d=1): # kernel, padding, dilation
if d > 1:
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size
if p is None:
p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
return p
class Conv(nn.Module):
default_act = nn.SiLU() # default activation
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x): # (b,c1,...)
# (b,c1,...)-->(b,c2,...),卷积块
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
return self.act(self.conv(x))
class Conv2(Conv):
def __init__(self, c1, c2, k=3, s=1, p=None, g=1, d=1, act=True):
super().__init__(c1, c2, k, s, p, g=g, d=d, act=act)
self.cv2 = nn.Conv2d(c1, c2, 1, s, autopad(1, p, d), groups=g, dilation=d, bias=False) # add 1x1 conv
def forward(self, x): # (b,c1,...)
# self.cv2(x):(b,c1,...)-->(b,c2,...) # 点卷积
# self.conv(x):(b,c1,...)-->(b,c2,...) # 普通卷积
# 先把两种卷积的处理结果做残差,之后批次标准化,激活函数
return self.act(self.bn(self.conv(x) + self.cv2(x)))
def forward_fuse(self, x): # (b,c1,...)
return self.act(self.bn(self.conv(x)))
def fuse_convs(self):
# 具有conv权重w形状的全0张量
w = torch.zeros_like(self.conv.weight.data)
i = [x // 2 for x in w.shape[2:]] # (1,1)
# 将w中最后两维1:2的数据用cv2的权重替换
w[:, :, i[0] : i[0] + 1, i[1] : i[1] + 1] = self.cv2.weight.data.clone()
# 用conv.weight和w做残差,结果做为conv的权重
self.conv.weight.data += w
self.__delattr__("cv2") # 删除cv2属性
self.forward = self.forward_fuse # 更改对象的forward方法
class DFL(nn.Module):
def __init__(self, c1=16):
super().__init__()
# 点卷积
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
# 设置conv的权重
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1)) # (1,c,1,1)
self.c1 = c1
def forward(self, x):
b, _, a = x.shape # batch, channels, anchors
# (b,c,a)-->(b,4,16,a)-->(b,16,4,a)-->softmax-->conv-->(b,4,a)
# softmax在索引1的轴归一化
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
class Proto(nn.Module):
def __init__(self, c1, c_=256, c2=32):
super().__init__()
self.cv1 = Conv(c1, c_, k=3) # 普通卷积,核大小3
# 上采样,核大小是2
self.upsample = nn.ConvTranspose2d(c_, c_, 2, 2, 0, bias=True)
self.cv2 = Conv(c_, c_, k=3) # 普通卷积,核大小3
self.cv3 = Conv(c_, c2)# 点卷积
def forward(self, x):
# (b,h,w,c1)-->(b,h,w,c_)-->(b,2h,2w,c_)-->(b,2h,2w,c_)-->(b,2h,2w,c2)
return self.cv3(self.cv2(self.upsample(self.cv1(x))))
class HGStem(nn.Module):
def __init__(self, c1, cm, c2):
super().__init__()
# 核大小3,步长2
self.stem1 = Conv(c1, cm, 3, 2, act=nn.ReLU())
# 核大小2,步长1
self.stem2a = Conv(cm, cm // 2, 2, 1, 0, act=nn.ReLU())
self.stem2b = Conv(cm // 2, cm, 2, 1, 0, act=nn.ReLU()) # 核大小2,步长1
self.stem3 = Conv(cm * 2, cm, 3, 2, act=nn.ReLU()) # 核大小3,步长2
self.stem4 = Conv(cm, c2, 1, 1, act=nn.ReLU()) # 点卷积
# ceil_mode=True 指的是当计算输出尺寸的时候，使用向上取整而不是向下取整（默认行为）。这在某些情
# 况下可能会导致输出的大小比通常预期的大一点。
# 举个例子，如果输入的尺寸为 (N, C, H, W)，其中 N 是batch size，C 是通道数，H 和 W 分别是高
# 度和宽度。假设 H 和 W 都是偶数，比如 (4, 4)，那么应用上述的 MaxPool2d 层之后，如果使用默认的
# floor 模式，输出的高度和宽度将是 (3, 3)；但如果使用了 ceil_mode=True，输出的高度和宽度将会是 (4, 4)。
self.pool = nn.MaxPool2d(kernel_size=2, stride=1, padding=0, ceil_mode=True)
def forward(self, x): # (b,h,w,c1)
x = self.stem1(x) # (b,h/2,w/2,cm)
x = F.pad(x, [0, 1, 0, 1])
x2 = self.stem2a(x) # (b,h/2,w/2,cm/2)
x2 = F.pad(x2, [0, 1, 0, 1])
x2 = self.stem2b(x2) # (b,h/2,w/2,cm)
x1 = self.pool(x) # (b,h/2,w/2,cm)
x = torch.cat([x1, x2], dim=1) # (b,h/2,w/2,2cm)
x = self.stem3(x) # (b,h/4,w/4,cm)
x = self.stem4(x) # (b,h/4,w/4,c2)
return x
class HGBlock(nn.Module):
def __init__(self, c1, cm, c2, k=3, n=6, lightconv=False, shortcut=False, act=nn.ReLU()):
super().__init__()
block = LightConv if lightconv else Conv
# 在nn.ModuleList中存储的对象是不同的对象，它们在内存中的地址也是不同的。nn.ModuleList是一个有序集合，
# 用于存储子模块，它可以像普通的Python列表一样索引，但它确保每个添加的项都是torch.nn.Module的实例或其子类。
# 每次调用block()时，都会创建一个新的实例，并将其添加到ModuleList中。这意味着每个block对象都有自己的状态和参数，
# 并且在模型的前向传播过程中会被独立地调用。
self.m = nn.ModuleList(block(c1 if i == 0 else cm, cm, k=k, act=act) for i in range(n))
self.sc = Conv(c1 + n * cm, c2 // 2, 1, 1, act=act) # 点卷积
self.ec = Conv(c2 // 2, c2, 1, 1, act=act) # 点卷积
self.add = shortcut and c1 == c2 # 残差条件
def forward(self, x):
y = [x]
# 这些block对象依次应用于输入数据
# 这段代码遍历self.m中的每个block对象，并将上一个block的输出作为当前block的输入，从而形成一个序
# 列化的处理流程。因为每次m(y[-1])处理后的结果都被添加到y中,而y[-1]是取出列表最后一个元素
# 这样每次m处理后的结果都被当成了下次的输入
y.extend(m(y[-1]) for m in self.m) # 6个卷积块处理
# 在通道轴合并特征
y = self.ec(self.sc(torch.cat(y, 1)))
# 满足残差条件,返回残差,否则返回处理后结果
return y + x if self.add else y
class SPP(nn.Module):
def __init__(self, c1, c2, k=(5, 9, 13)):
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1) # 点卷积
self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) # 点卷积
self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
def forward(self, x): # (b,h,w,c1)
x = self.cv1(x) # (b,h,w,c_)
# (b,h,w,c_)-->(b,h,w,4c_)-->(b,h,w,c2)
# 这里不同尺寸的最大池化核会处理相同的x,之后列表内的对象在索引1的轴(通道轴)合并
# self.m中的每个m处理的都是相同的输入x，而不是上一个最大池化的输出
return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
class SPPF(nn.Module):
def __init__(self, c1, c2, k=5):
super().__init__()
c_ = c1 // 2 # hidden channels
self.cv1 = Conv(c1, c_, 1, 1) # 点卷积
self.cv2 = Conv(c_ * 4, c2, 1, 1) # 点卷积
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
# (b,h,w,c1)-->(b,h,w,c_)
y = [self.cv1(x)]
# 通过列表推导式和extend方法将多次迭代的结果追加到y列表中。每次迭代都会使用self.m
# 处理列表y的最后一个元素，并将结果追加到y列表中。因此，y[-1]始终是指向列表中最新追加的元素
y.extend(self.m(y[-1]) for _ in range(3))
# 将列表y中的所有元素沿着通道维度（维度1）拼接起来
# (b,4c_,h,w)-->(b,c2,h,w)
return self.cv2(torch.cat(y, 1))
class C1(nn.Module):
def __init__(self, c1, c2, n=1):
super().__init__()
self.cv1 = Conv(c1, c2, 1, 1) # 点卷积
self.m = nn.Sequential(*(Conv(c2, c2, 3) for _ in range(n)))
def forward(self, x):
# (b,c1,h,w)-->(b,c2,h,w)
y = self.cv1(x)
# 序列化栈前后残差
return self.m(y) + y
class Bottleneck(nn.Module): # 标准瓶颈模块
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, k[0], 1) # 普通卷积
self.cv2 = Conv(c_, c2, k[1], 1, g=g) # g是groups
# 残差条件:shortcut为True并且c1 == c2
self.add = shortcut and c1 == c2
def forward(self, x): # (b,c1,h,w)
# self.cv2(self.cv1(x)):(b,c1,h,w)-->(b,c_,h,w)-->(b,c2,h,w)
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
class C2(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__()
self.c = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, 2 * self.c, 1, 1) # 点卷积
self.cv2 = Conv(2 * self.c, c2, 1)
self.m = nn.Sequential(*(
Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)))
def forward(self, x): # (b,c1,h,w)
# (b,c1,h,w)-->(b,2c,h,w)
# chunk(2, 1)：将输出张量沿着通道维度（维度1）分割成两个部分，分别命名为a和b。
# 2是指分成两个部分,1指索引1的轴
a, b = self.cv1(x).chunk(2, 1)
# 将分割出的a传递给模块self.m,将处理后的a与未处理的b沿通道维度拼接起来。
# 将拼接后的张量再次通过卷积层cv2进行处理
# (b,2c,...)-->(b,c2,...)
return self.cv2(torch.cat((self.m(a), b), 1))
class C2f(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
super().__init__()
self.c = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2)
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x):# (b,c1,...)
# (b,c1,...)-->(b,2c,...),在通道维度分成两份
y = list(self.cv1(x).chunk(2, 1))
# 依次追加self.m中每个m处理后的结果,y[-1]是上个处理结果
y.extend(m(y[-1]) for m in self.m)
# 将列表中的所有特征图在通道轴合并
# (b,(2 + n)c,...)-->(b,c2,...)
return self.cv2(torch.cat(y, 1))
def forward_split(self, x):
y = list(self.cv1(x).split((self.c, self.c), 1))
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
class SCDown(nn.Module):
def __init__(self, c1, c2, k, s):
super().__init__()
self.cv1 = Conv(c1, c2, 1, 1) # 点卷积
self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False)
def forward(self, x): # (b,c1,...)
# (b,c1,...)-->(b,c2,...)-->(b,c2,...)
# 先点卷积切换通道,之后深度卷积
return self.cv2(self.cv1(x))
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, attn_ratio=0.5):
super().__init__()
self.num_heads = num_heads # h
self.head_dim = dim // num_heads # dk
self.key_dim = int(self.head_dim * attn_ratio)
self.scale = self.key_dim**-0.5 # 缩放系数
nh_kd = self.key_dim * num_heads # 所有头加起来的维数
h = dim + nh_kd * 2 # 中间维度
self.qkv = Conv(dim, h, 1, act=False)
self.proj = Conv(dim, dim, 1, act=False) # 点卷积
# 深度卷积
self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)
def forward(self, x): # (b,d,...)
B, C, H, W = x.shape
N = H * W
# (b,d,...)-->(b,h,...)
qkv = self.qkv(x)
# (b,h,...)-->(b,h,dk+k_dk*2,s),在通道维度进行拆分,q,k特征维度相同
# v单独一个特征维度
# q和k主要用于计算注意力权重，因此减少它们的特征维度可以减少冗余信息，提高计算效率。
# v用于根据注意力权重进行加权求和，保持较大的特征维度可以保留更多的信息。
q, k, v = qkv.view(B, self.num_heads, self.key_dim * 2 + self.head_dim, N).split(
[self.key_dim, self.key_dim, self.head_dim], dim=2
)
# (b,h,s_q,k_dk)@(b,h,k_dk,s_k)-->(b,h,s_q,s_k)
attn = (q.transpose(-2, -1) @ k) * self.scale
attn = attn.softmax(dim=-1) # 在s_k上归一化
# (b,h,dk,s_v)@(b,h,s_k,s_q)-->(b,h,dk,s_q)-->(b,d,h,w)
# v:(b,h,dk,s)-->(b,d,h,w)-->(b,d,...)
# 注意力后的v和经过深度卷积的v做残差
x = (v @ attn.transpose(-2, -1)).view(B, C, H, W) + self.pe(v.reshape(B, C, H, W))
# (b,d,h,w)-->(b,d,h,w)
x = self.proj(x)
return x
class PSA(nn.Module):
def __init__(self, c1, c2, e=0.5):
super().__init__()
assert c1 == c2
self.c = int(c1 * e)
self.cv1 = Conv(c1, 2 * self.c, 1, 1) # 点卷积切通道
self.cv2 = Conv(2 * self.c, c1, 1)
self.attn = Attention(self.c, attn_ratio=0.5, num_heads=self.c // 64)
self.ffn = nn.Sequential(Conv(self.c, self.c * 2, 1), Conv(self.c * 2, self.c, 1, act=False))
def forward(self, x): # (b,c1,h,w)
# (b,c1,h,w)-->(b,2c,...),split会在通道维度拆分
a, b = self.cv1(x).split((self.c, self.c), dim=1)
# 对b做注意力前后的残差连接
b = b + self.attn(b) # (b,c,...)
# (b,c,...)-->(b,2c,...)-->(b,c,...)
# 整个多头自注意力过程没有层标准化和dropout
b = b + self.ffn(b) # 前馈前后残差
# 在通道维度合并a,b的特征,之后经过卷积处理
# (b,2c,...)-->(b,c1,...)
return self.cv2(torch.cat((a, b), 1))
class C2fCIB(C2f):
def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
self.m = nn.ModuleList(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))
class C3(nn.Module):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__()
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) # optional act=FReLU(c2)
self.m = nn.Sequential(*(Bottleneck(c_, c_, shortcut, g, k=((1, 1), (3, 3)), e=1.0) for _ in range(n)))
def forward(self, x): # (b,c1,...)
# x经过cv1,之后被m模块处理,另一个x经过cv2处理,之后两者在通道维度合并特征
# (b,c1,...)-->(b,c_,...)-->(b,c,...)
# b,c1,...)-->(b,c_,...)
# 之后在通道维度合并,-->(b,2c,...)-->(b,c2,...)
return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
class C3x(C3):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
self.c_ = int(c2 * e)
self.m = nn.Sequential(*(Bottleneck(self.c_, self.c_, shortcut, g, k=((1, 3), (3, 1)), e=1) for _ in range(n)))
class C3TR(C3):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e)
self.m = TransformerBlock(c_, c_, 4, n)
class C3Ghost(C3):
def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
super().__init__(c1, c2, n, shortcut, g, e)
c_ = int(c2 * e) # hidden channels
self.m = nn.Sequential(*(GhostBottleneck(c_, c_) for _ in range(n)))
class RepC3(nn.Module):
def __init__(self, c1, c2, n=3, e=1.0):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c2, 1, 1)
self.cv2 = Conv(c1, c2, 1, 1)
self.m = nn.Sequential(*[RepConv(c_, c_) for _ in range(n)])
self.cv3 = Conv(c_, c2, 1, 1) if c_ != c2 else nn.Identity()
def forward(self, x):
return self.cv3(self.m(self.cv1(x)) + self.cv2(x))