YOLO v5-C3模块实现

小啊磊_Vv

已于 2023-11-10 17:25:43 修改

阅读量3.4k

点赞数 5

文章标签：深度学习 YOLO 人工智能计算机视觉神经网络

于 2023-03-10 11:31:47 首次发布

本文链接：https://blog.csdn.net/m0_60890175/article/details/129436954

版权

YOLO v5-C3模块实现

文章目录

YOLO v5-C3模块实现
前言
一、定义自动padding
二、基本Conv模块
三、Bottleneck模块
四、C3模块
五、其他模块
六、编写模块构建网络实现分类

前言

总结

在本周的学习中自己学会YOLO v5中的C3模块的搭建, 同时根据网络结构图,实现SPP, SPPF 等模块, 并用几个简单的模块,对之前的数据集实现分类
自己在学习过程是，知道class一些实现方法，写网络要找到共同的地方，定义基本模块，实现模块的复用。
如何确定第一个linear层的神经元个数

一、定义自动padding

为了保持图像大小卷积前后一致,就需要用到自动padding

def autopad(k, p=None):                        # kernel  padding 根据卷积核大小k自动计算卷积核padding数（0填充）
    """
    :param k: 卷积核的 kernel_size
    :param p: 卷积的padding  一般是None
    :return:  自动计算的需要pad值(0填充)
    """
    if p is None:
        # k 是 int 整数则除以2, 若干的整数值则循环整除
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
    return p

二、基本Conv模块

C3, SPP, SPPF模块中都有基本的Conv模块,需要编写基本的Conv模块,方面复用，提高代码的复用性

class Conv(nn.Module):
    def __init__(self, c1, c2, k=1, s=1, p=None, act=True, g=1):
        """
        :param c1: 输入的channel值
        :param c2: 输出的channel值
        :param k: 卷积的kernel_size
        :param s: 卷积的stride
        :param p: 卷积的padding  一般是None
        :param act: 激活函数类型   True就是SiLU(), False就是不使用激活函数
        :param g: 卷积的groups数  =1就是普通的卷积  >1就是深度可分离卷积
        """
        super(Conv, self).__init__()

        self.conv_1 = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=True)
        self.bn = nn.BatchNorm2d(c2)

        self.act = nn.SiLU() if act else nn.Identity()     # 若act=True, 则激活,  act=False, 不激活

    def forward(self, x):

        return self.act(self.bn(self.conv_1(x)))

三、Bottleneck模块

Bottleneck模块中包含一个残差连接结构（左），和不包含的残差结构（右），就需要传入参数,来判断是否需要使用残差结构

class Bottleneck(nn.Module):
    def __init__(self, c1, c2, e=0.5, shortcut=True, g=1):
        """
        :param c1: 整个Bottleneck的输入channel
        :param c2: 整个Bottleneck的输出channel
        :param e: expansion ratio  c2*e 就是第一个卷积的输出channel=第二个卷积的输入channel
        :param shortcut: bool Bottleneck中是否有shortcut，默认True
        :param g: Bottleneck中的3x3卷积类型  =1普通卷积  >1深度可分离卷积
        """
        super(Bottleneck, self).__init__()

        c_ = int(c2*e)                            # 使通道减半, c_具体多少取决于e
        self.conv_1 = Conv(c1, c_, 1, 1)
        self.conv_2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.conv_2(self.conv_1(x)) if self.add else self.conv_2(self.conv_1(x))

四、C3模块

class C3(nn.Module):
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        """
        :param c1: 整个 C3 的输入channel
        :param c2: 整个 C3 的输出channel
        :param n: 有n个Bottleneck
        :param shortcut: bool Bottleneck中是否有shortcut，默认True
        :param g: C3中的3x3卷积类型  =1普通卷积  >1深度可分离卷积
        :param e: expansion ratio
        """
        super(C3, self).__init__()
        c_ = int(c2 * e)
        self.cv_1 = Conv(c1, c_, 1, 1)
        self.cv_2 = Conv(c1, c_, 1, 1)
        # *操作符可以把一个list拆开成一个个独立的元素,然后再送入Sequential来构造m,相当于m用了n次Bottleneck的操作
        self.m = nn.Sequential(*[Bottleneck(c_, c_, e=1, shortcut=True, g=1) for _ in range(n)])
        self.cv_3 = Conv(2*c_, c2, 1, 1)

    def forward(self, x):
        return self.cv_3(torch.cat((self.m(self.cv_1(x)), self.cv_2(x)), dim=1))

五、其他模块

class BottleneckCSP(nn.Module):
    def __init__(self, c1, c2, e=0.5, n=1):
        """
        :param c1: 整个BottleneckCSP的输入channel
        :param c2: 整个BottleneckCSP的输出channel
        :param e: expansion ratio c2*e=中间其他所有层的卷积核个数/中间所有层的输入输出channel数
        :param n: 有 n 个Bottleneck
        """
        super(BottleneckCSP, self).__init__()
        c_ = int(c2*e)
        self.conv_1 = Conv(c1, c_, 1, 1)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, e=1, shortcut=True, g=1) for _ in range(n)])
        self.conv_3 = Conv(c_, c_, 1, 1)
        self.conv_2 = Conv(c1, c_, 1, 1)
        self.bn = nn.BatchNorm2d(2*c_)
        self.LeakyRelu = nn.LeakyReLU()
        self.conv_4 = Conv(2*c_, c2, 1, 1)

    def forward(self, x):
        x_1 = self.conv_3(self.m(self.conv_1(x)))
        x_2 = self.conv_2(x)
        x_3 = torch.cat([x_1, x_2], dim=1)
        x_4 = self.LeakyRelu(self.bn(x_3))
        x = self.conv_4(x_4)
        return x

class SPP(nn.Module):
    def __init__(self, c1, c2, e=0.5, k1=5, k2=9, k3=13):
        """
        :param c1: SPP模块的输入channel
        :param c2: SPP模块的输出channel
        :param e: expansion ratio
        :param k1: Maxpool 的卷积核大小
        :param k2: Maxpool 的卷积核大小
        :param k3: Maxpool 的卷积核大小
        """
        super(SPP, self).__init__()
        c_ = int(c2*e)
        self.cv_1 = Conv(c1, c_, 1, 1)
        self.pool_1 = nn.MaxPool2d(kernel_size=k1, stride=1, padding=k1 // 2)
        self.pool_2 = nn.MaxPool2d(kernel_size=k2, stride=1, padding=k2 // 2)
        self.pool_3 = nn.MaxPool2d(kernel_size=k3, stride=1, padding=k3 // 2)
        self.cv_2 = Conv(4*c_, c2, 1, 1)

    def forward(self, x):
        return self.cv_2(torch.cat((self.pool_1(self.cv_1(x)), self.pool_2(self.cv_1(x)), self.pool_3(self.cv_1(x)), self.cv_1(x)), dim=1))

六、编写模块构建网络实现分类

6.1 数据集操作

total_dir = './weather_photos/'

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(std=[0.5, 0.5, 0.5], mean=[0.5, 0.5, 0.5])
])

total_data = torchvision.datasets.ImageFolder(total_dir, transform)
print(total_data)
print(total_data.class_to_idx)

idx_to_class = dict((v, k) for k,v in total_data.class_to_idx.items())
print(idx_to_class)

train_size = int(len(total_data) * 0.8)
test_size = int(len(total_data)) - train_size

train_dataset, test_dataset = torch.utils.data.random_split(total_data, [train_size, test_size])

train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=32, shuffle=True)

6.1 构建网络

在编写上述模块时，发现最主要的参数就是输入通道，和输入通道，只需要保证调用模块时，上次一层的输入通道与下一层输出通道保持一致就行。

6.2 linear神经元个数判定

在构建网络模块时不知道第一个全连接神经元个数，可用如下方法进行判断，构建好初全连接以外的层，在forward进行传播,并在最后一个层前向传播完进行打印size。

class model(nn.Module):
	def __init__(self):
        super(model, self).__init__()

        self.conv = Conv(3, 32, 3, 2)                   # 3:输入通道  32:输出通道  3: kernel  2:stride
        self.spp = SPP(32, 64)
        self.c3 = C3(64, 128, n=1, shortcut=True, g=1, e=0.5)
        # 接下来就是 linear 层

    def forward(self, x):
        x = self.conv(x)
        x = self.spp(x)
        x = self.c3(x)
        print(x.size())
        return x


device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model().to(device)

for x, y in train_dataloader:
    x, y = x.to(device), y.to(device)
    y_pre = model(x)
    break

打印输出 torch.Size([32, 128, 112, 112])， 32为batch-size，全连接神经元个数就是128112112，
重新编写网络结构函数

class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()

        self.conv = Conv(3, 32, 3, 2)                   # 3:输入通道  32:输出通道  3: kernel  2:stride
        self.spp = SPP(32, 64)
        self.c3 = C3(64, 128, n=1, shortcut=True, g=1, e=0.5)
        self.linear = nn.Sequential(
            nn.Linear(128*112*112, 1000),
            nn.ReLU(),

            nn.Linear(1000, 4)
        )

    def forward(self, x):
        x = self.conv(x)
        x = self.spp(x)
        x = self.c3(x) 
        x = x.view(-1, 128*112*112)
        x = self.linear(x)
        return x