YOLOv5中SPP/SPPF结构源码详析(内含注释分析)

class SPP(nn.Module):
    def __init__(self, c1, c2, k=(5, 9, 13)):#这里5，9，13，就是初始化的kernel size
        super().__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)#这里对应第一个CBL
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)#这里对应SPP操作里的最后一个CBL
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
        #这里对应SPP核心操作，对5 * 5分块，9 * 9分块和13 * 13子图分别取最大池化
 
    def forward(self, x):
        x = self.cv1(x)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')  # suppress torch 1.9.0 max_pool2d() warning忽略警告
            return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
            #torch.cat对应concat

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)#先通过CBL进行通道数的减半
        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))
            #将原来的x,一次池化后的y1,两次池化后的y2,3次池化的self.m(y2)先进行拼接，然后再CBL