YOLO v5-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
- 并运行代码: Process finished with exit code 0
- 按照之前的代码进行训练评估即可