- class Conv(nn.Module):
- # Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)
- 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):
- return self.act(self.bn(self.conv(x)))
-
- def forward_fuse(self, x):
- return self.act(self.conv(x))
代码定义了一个名为Conv的类，它是nn.Module的子类。该类执行一个标准的卷积操作，可选地包括激活函数和批归一化。
该类具有以下属性和方法：
default_act：这是一个类属性，设置为nn.SiLU()（Sigmoid-weighted Linear Unit），它是卷积操作中使用的默认激活函数。__init__：这是构造方法，用于初始化Conv对象。它接受以下参数：c1：输入通道数。c2：输出通道数。k：卷积核大小（默认为1）。s：步长大小（默认为1）。p：填充大小（默认为None）。g：分组数（默认为1）。d：膨胀因子（默认为1）。act：激活函数或True（默认为True）。在这个方法中，初始化了一个Conv2d层、BatchNorm2d层和激活函数，并将它们存储为属性。
在该方法中，首先计算隐藏通道数c_，然后初始化了三个Conv层 cv1, cv2 和 cv3，其中 cv1 和 cv2 是1x1的卷积层，cv3 是具有输出通道数c2的1x1的卷积层。这些Conv层用于处理输入数据。 然后，使用nn.Sequential()初始化了一个包含n个Bottleneck模块的Sequential网络，每个Bottleneck模块的输入和输出通道数都是c_。
在该方法中，首先计算隐藏通道数c_，然后初始化了两个Conv层 cv1 和 cv2。其中 cv1 是用于处理输入张量的1x1卷积层，将输入通道数减少到c_。cv2 是用于最终输出的1x1卷积层，将输入通道数增加到c2。 接下来，使用nn.MaxPool2d初始化了一个最大池化层，池化核大小为k，步长为1，填充大小为k // 2。这个最大池化层用于执行空间金字塔池化操作。
forward：该方法执行Conv对象的前向传播。它接受一个输入张量x，按照顺序应用卷积操作、批归一化和激活函数。然后返回输出张量。
forward_fuse：该方法类似于forward，但不包括批归一化。它只执行卷积操作和激活函数。然后返回输出张量。
- class C3(nn.Module):
- # CSP Bottleneck with 3 convolutions
- def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
- 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, e=1.0) for _ in range(n)))
-
- def forward(self, x):
- return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), 1))
代码定义了一个名为C3的类，它是nn.Module的子类。该类实现了一个带有3个卷积层的CSP Bottleneck结构。
该类有以下属性和方法：
__init__：构造方法，用于初始化C3对象。它接受以下参数：c1：输入通道数。forward：该方法执行C3对象的前向传播。它接受一个输入张量x，首先将输入张量通过cv1进行卷积，然后将输入张量通过cv2进行卷积，并将两者结果进行拼接。接着将拼接后的张量作为输入传递给Sequential网络m进行处理。最后，将处理后的结果和通过cv3进行卷积的输入张量进行拼接，并返回最终的输出张量。
c2：输出通道数。n：重复次数（默认为1）。shortcut：是否使用残差连接（默认为True）。g：分组数（默认为1）。e：扩展系数（默认为0.5）。 - class SPPF(nn.Module):
- # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
- def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13))
- 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):
- x = self.cv1(x)
- with warnings.catch_warnings():
- warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning
- y1 = self.m(x)
- y2 = self.m(y1)
- return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
码定义了一个名为SPPF的类，它是nn.Module的子类。该类实现了一个用于YOLOv5的SPPF（Spatial Pyramid Pooling - Fast）层。
该类有以下属性和方法：
__init__：构造方法，用于初始化SPPF对象。它接受以下参数：c1：输入通道数。forward：该方法执行SPPF对象的前向传播。它接受一个输入张量x，首先通过cv1对输入张量进行1x1的卷积。接下来，使用最大池化层m对cv1的输出张量进行多次池化操作，生成不同尺度的特征图。最后，将cv1的输出张量、池化特征图进行拼接，并通过cv2进行1x1的卷积，得到最终的输出张量。
c2：输出通道数。k：池化核大小（默认为5）。 - def parse_model(d, ch): # model_dict, input_channels(3)
- # Parse a YOLOv5 model.yaml dictionary
- LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
- anchors, nc, gd, gw, act = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation')
- if act:
- Conv.default_act = eval(act) # redefine default activation, i.e. Conv.default_act = nn.SiLU()
- LOGGER.info(f"{colorstr('activation:')} {act}") # print
- na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
- no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
-
- layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
- for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
- m = eval(m) if isinstance(m, str) else m # eval strings
- for j, a in enumerate(args):
- with contextlib.suppress(NameError):
- args[j] = eval(a) if isinstance(a, str) else a # eval strings
-
- n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
- if m in {
- Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
- BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
- c1, c2 = ch[f], args[0]
- if c2 != no: # if not output
- c2 = make_divisible(c2 * gw, 8)
-
- args = [c1, c2, *args[1:]]
- if m in {BottleneckCSP, C3, C3TR, C3Ghost, C3x}:
- args.insert(2, n) # number of repeats
- n = 1
- elif m is nn.BatchNorm2d:
- args = [ch[f]]
- elif m is Concat:
- c2 = sum(ch[x] for x in f)
- # TODO: channel, gw, gd
- elif m in {Detect, Segment}:
- args.append([ch[x] for x in f])
- if isinstance(args[1], int): # number of anchors
- args[1] = [list(range(args[1] * 2))] * len(f)
- if m is Segment:
- args[3] = make_divisible(args[3] * gw, 8)
- elif m is Contract:
- c2 = ch[f] * args[0] ** 2
- elif m is Expand:
- c2 = ch[f] // args[0] ** 2
- else:
- c2 = ch[f]
-
- m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module
- t = str(m)[8:-2].replace('__main__.', '') # module type
- np = sum(x.numel() for x in m_.parameters()) # number params
- m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
- LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}') # print
- save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
- layers.append(m_)
- if i == 0:
- ch = []
- ch.append(c2)
- return nn.Sequential(*layers), sorted(save)
代码定义了一个名为parse_model的函数，用于解析YOLOv5的model.yaml字典。
该函数接受两个参数：d是一个包含YOLOv5模型配置信息的字典，ch是输入通道数的列表。
函数的主要功能是根据模型配置信息解析出网络的结构，并返回一个包含所有层的nn.Sequential对象以及一个已排序的保存列表。
具体的解析过程如下：