文章目录[隐藏]
Backbone
激活函数的定义–Mish
import math
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F
from model.layers.attention_layers import SEModule, CBAM
import config.yolov4_config as cfg
class Mish(nn.Module):
def __init__(self):
super(Mish, self).__init__()
def forward(self, x):
return x*torch.tanh(F.softplus(x))
这一块就是定义了一个我们在yolov4中即将用到的一个激活函数Mish(),这个激活函数将出现在每个卷积模块中。
Mish激活函数的优点:
以上无边界(即正值可以达到任何高度)避免了由于封顶而导致的饱和。理论上对负值的轻微允许允许更好的梯度流,而不是像ReLU中那样的硬零边界,而且平滑的激活函数允许更好的信息深入神经网络,从而得到更好的准确性和泛化。
各类激活函数可以参考大神bubbliiiing的这篇博客:各类激活函数Activation Functions介绍与优缺点分析
全局变量的定义
norm_name = {"bn":nn.BatchNorm2d}
activate_name = {
"relu":nn.ReLU,
"leaky":nn.LeakyReLU,
"linear":nn.Identity,
"mish":Mish(),
}
这里主要定义了全局变量,通过字典的形式,方便我们在接下去写代码的时候调用各种torch.nn的各种工具函数,增加了代码的可读性。
卷积模块的定义–CBM
class Convolutional(nn.Module):
def __init__(
self,
filters_in,
filters_out,
kernel_size,
stride=1,
norm="bn",
activate="mish",
):
super(Convolutional, self).__init__()
self.norm = norm
self.activate = activate
self.__conv = nn.Conv2d(
in_channels = filters_in,
out_channels = filters_out,
kernel_size = kernel_size,
stride = stride,
padding = kernel_size//2,
bias = not norm,
)
if norm:
assert norm in norm_name.keys
if norm == "bn":
self.__norm = norm_name[norm](num_features=filters_out)
if activate:
assert activate in activate_name.keys()
if activate == "leaky":
self.__activate = activate = activate_name[activate](
negative_slope = 0.1, inplace=True
)
if activate == "relu":
self.__activate = activate_name[activate](inplace=True)
if activate == "mish":
self._-activate = activate_name[activate]
def forward(self, x):
x = self.__conv(x)
if self.norm:
x = self.__norm(x)
if self.activate:
x = self.__activate(x)
return x
在这部分,我们主要完成了一个CBM卷积模块的一个定义,其中涉及到一次卷积运算,一次BatchNorm运算和一次Mish激活函数运算。顺序正如这段代码所示,在前向函数中,对形参x先是卷积再bn算法再Mish激活。
代码中的padding=kernel_size//2是一种向下取整的方式,目的是为了保持在不同卷积核尺寸下得到的特征图的大小一致(ps:padding的这种取值方式的前提是stride=1)
代码中涉及到了三种激活函数,包括leaky,relu,以及mish,在YOLOv4中我们使用的是mish。
代码中的bias=not norm应该是指当采用BN算法时,不用进行偏置运算,如果不刻意设置,默认为true。在YOLOv4中我们采取的是BN算法,因此norm设置为ture,if判断语句会选择将BN算法赋给norm,以便后面的调用。可视化如图:
小残差模块的定义–Resunit
class CSPBlock(nn.Module):
def __init__(
self,
in_channels,
out_channels,
hidden_channels = None,
residual_activation = "linear",
):
super(CSPBlock, self).__init__()
if hidden_channels is None:
hidden_channels = out_channels
self.block = nn.Sequential(
Convolutional(in_channels,hidden_channels, 1),
Convolutional(hidden_channels,out_channels, 3),
)
self.activation = activate_name[residual_activation]
self.attention = cfg.ATTENTION["TYPE"]
if self.attention == "SEnet":
self.attention_module = SEModule(out_channels)
elif self.attention == "CBAM":
self.attention_module = CBAM(out_channels)
elif
self.attention == None
def forward(self, x):
residual = x
out = self.block(x)
if self.attention is not None:
out = self.attention_module(out)
out += residual
return out
在这一部分代码中,定义了残差模块Resunit,该模块由两个CBM小模块和一个残差边组成。其中两个CBM模块构成的小网络定义在self.block中,其中一个卷积核尺寸为1,一个卷积核尺寸为3。而后定义的一个self.activation在代码中并没有调用,我的理解是此处将残差边的激活函数赋给self.activation,而该激活函数是key为“linear”的激活函数,找到上面定义的全局变量可知,对应的激活函数为nn.Identity(),该激活函数通过查询了解到,其在网络中的作用仅仅是增加了层数,对我们的输入并没有其他的操作,可以理解为是一个桥的作用,因此key的名字为linear,即为线性映射。由于并没有实质性的作用,因此作者在下面的代码中并未出现调用的地方(也许调用了,只是我没有发现~~doge)。讲到哪了来着,现在我们定义好了Resunit的卷积边,然后下面根据三种注意力算法,定义了三种使用情况:SEnet,CBAM和None(具体采用哪种注意力机制,取决于配置文件中的设置,在config文件中)。具体注意力机制在YOLO算法中的作用,可以查看这篇博客:SEnet,CBAM。总之,简单概况,注意力机制能够有效的提高图像分类和目标检测的准确率。继续往下看,来到了前向函数部分,这一部分可以清楚的看到,输入经过了我们上述定义的self.block卷积网络和注意力算法(如果有的话),得到输出out,然后将我们的输入定义为residual(中文译为残余强度,在YOLO中即为残差边),add在前面得到的out上,得到最终的输出。到这里,我们的残差模块Resunit就定义好了~。网络可视化如图:
大残差模块的定义–CSP1
class CSPFirstStage(nn.Module):
def __init__(self, in_channels, out_channels):
super(CSPFirstStage, self).__init__()
self.downsample_conv = Convolutional(in_channels, out_channels, 3, stride=2)
self.split_conv0 = Convolutional(out_channels, out_channels, 1)
self.split_conv1 = Convolutional(out_channels, out_channels, 1)
self.blocks_conv = nn.Sequential(
CSPBlock(out_channels, out_channels, in_channels),
Convolutiona(out_channels, out_channels, 1),
)
self.concat_conv = Convolutional(out_channels * 2, out_channels, 1)
def forward(self, x):
x = self.downsample_conv(x)
x0 = self.split_conv0(x)
x1 = self.split_conv1(x)
x1 = self.block_conv(x1)
x = torch.cat([x0, x1], dim = 1)
x = self.concat_conv(x)
return x
在这一部分,我们定义了一个大残差模块。前文所定义的小残差模块将作为这个大残差模块的重要一部分(具体的可以查看YOLOv4网络可视化之后的框架图)。现在,开始讲解这部分的代码~。首先根据YOLOv4的网络框架图我们可以清晰的发现,每当我们的输入经过一次大残差模块时,都会被降采样一次,即输出特征图尺寸变为输入特征图的1/2,因此首先定义了一个降采样的CBM,然后为了区分输入,我们定义了两个CBM,输出分别通往不同的分支,一个通往小残差模块处,一个通往残差边。小残差模块这条线还存在一个小残差模块CSP和一个卷积模块CBM,因此定义了self.blocks_conv来构建这个网络。
然后根据前向函数可以看到,首先在主干上,也是输入的必经之路上存在一个降采样的CBM卷积模块,然后输出分成两个支线,一个经过存在小残差模块网络结构,一个经过大残差模块的残差边,最后根据self.concat_conv函数将两个分支的输出在维度上进行叠加。这部分定义的大残差网络是CSPDarknet53网络的第一块,也就是只有一个Resunit组件的CSP。这部分代码可视化如图:
大残差模块的定义-CSPx
class CSPStage(nn.Module):
def __init__(self, in_channels, out_channels, num_blocks):
super(CSPStage, self).__init__()
self.downsample_conv = Convolutional(
in_channels, out_channels, 3, stride = 2
)
self.split_conv0 = Convolutional(out_channels, out_channels//2, 1)
self.split_conv1 = Convolutional(out_channels, out_channels//2, 1)
self.blocks_conv = nn.Sequential(
*[
CSPBlock(out_channels//2 , out_channels//2)
for _ in range(num_blocks)
],
Convolutional(out_channels//2, out_channels//2, 1)
)
self.concat_conv = Convolutional(out_channels, out_channels, 1)
def forward(self, x):
x = self.downsample_conv(x)
x0 = self.split0_conv0(x)
x1 = self.split1_conv1(x)
x1 = self.blocks_conv(x1)
x = torch.cat([x0, x1], dim = 1)
x = self.concat_conv(x)
return x
这一部分的代码和之前的CSP1的代码基本类似,唯一的区别就是这里定义的大残差模块中所调用的Resunit组件的个数可以自定义,也就是这个类中多了一个变量,即num_blocks。还有一个区别就是,为了保证在concat之后,特征图的通道数保持不变,因此这里在最后堆叠之前,提前将前几次的卷积模块输出的通道数变成了out_channels//2,这样在最后堆叠的时候,输入与输出都是一个通道数,这点区别于CSPFirstStage中最后的(out_channels * 2, out_channels)。
CSPDarknet53网络搭建
以上就是YOLOv4主干网络我们所需要的所有模块的定义了,接下去就正式开始像搭积木似的将CSPDarknet53搭建起来吧!!
class CSPDarknet53(nn.Module):
def __init__(
self,
stem_channels = 32,
feature_channels = [64, 128, 256, 512, 1024],
num_features = 3,
weight_path = None,
resume = False,
):
super(CSPDarknet53, self).__init__()
self.stem_conv = Convolutional(3, stem_channels, 3)
self.stages = nn.ModuleList(
[
CSPFirststage(stem_channels, feature_channels[0]),
CSPStage(feature_channels[0], feature_channels[1], 2),
CSPStage(feature_channels[1], feature_channels[2], 8),
CSPStage(feature_channels[2], feature_channels[3], 8),
CSPStage(feature_channels[3], feature_channels[4], 4),
]
)
self.feature_channels = feature_channels
self.num_features = num_features
if weight_path and not resume:
self.load_CSPdarknet_weights(weight_path)
else:
self._initialize_weights()
def forward(self, x):
x = self.stem_conv(x)
features = []
for stage in self.stage:
x = stage(x)
features.append(x)
return feature[-self.num_features:]
在self.stages中,我们用nn.ModuleList()来构建我们的网络,其中CSPStage方法的第三个参数表示在大残差模块中,Resunit组件的个数。在forward中,建立了一个特征列表,并且利用一个循环,来遍历我们所构建的网络self.stages,这也是利用nn.ModuleList()方法来构建网络的好处。然后我们就可以将这五个csp组件所得到的的输出都append进这个空列表中,最后返回的是feature[-self.num_features:],而self.num_features已经定义好为3,。那么为什么只需要返回列表的最后三个特征输出呢,这里我们可以根据下图发现,输入进入网络后,真正传出到下一个网络结构的只有那三个输出,就是从csp8,csp8,csp4传出的输出,这也是我们返回列表后三个特征的原因。根据网络结构我们知道,输入首先会遇到一个CBM卷积模块,即self.stem_conv()。输入是3是因为一开始的输入图像是彩色图像,有rgb三通道的像素值,然后我们的输出需要变成32通道的特征图,做完这一步操作,只是通道数变多,特征图的尺寸依旧是原图尺寸(可见整个网络的可视化图)。接下去,每次经过一个csp模块时,都会进行一次降采样操作和一次特征图堆叠操作,因此对应的,输出特征图每次都是 输入特征图尺寸的1/2,而通道数都会变成输入的2倍。CSPDarknet53如图:
权重的初始化以及载入
def _initialize_weights(self):
print("**" * 10, "Initing CSPDarknet53 weights", "**" * 10)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2.0 / n))
if m.bias is not None:
m.bias.data.zero_()
print("initing {}".format(m))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
print("initing {}".format(m))
def load_CSPdarknet_weights(self, weight_file, cutoff=52):
"https://github.com/ultralytics/yolov3/blob/master/models.py"
print("load darknet weights : ", weight_file)
with open(weight_file, "rb") as f:
_ = np.fromfile(f, dtype=np.int32, count=5)
weights = np.fromfile(f, dtype=np.float32)
count = 0
ptr = 0
for m in self.modules():
if isinstance(m, Convolutional):
# only initing backbone conv's weights
# if count == cutoff:
# break
# count += 1
conv_layer = m._Convolutional__conv
if m.norm == "bn":
# Load BN bias, weights, running mean and running variance
bn_layer = m._Convolutional__norm
num_b = bn_layer.bias.numel() # Number of biases
# Bias
bn_b = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.bias.data
)
bn_layer.bias.data.copy_(bn_b)
ptr += num_b
# Weight
bn_w = torch.from_numpy(weights[ptr : ptr + num_b]).view_as(
bn_layer.weight.data
)
bn_layer.weight.data.copy_(bn_w)
ptr += num_b
# Running Mean
bn_rm = torch.from_numpy(
weights[ptr : ptr + num_b]
).view_as(bn_layer.running_mean)
bn_layer.running_mean.data.copy_(bn_rm)
ptr += num_b
# Running Var
bn_rv = torch.from_numpy(
weights[ptr : ptr + num_b]
).view_as(bn_layer.running_var)
bn_layer.running_var.data.copy_(bn_rv)
ptr += num_b
print("loading weight {}".format(bn_layer))
else:
# Load conv. bias
num_b = conv_layer.bias.numel()
conv_b = torch.from_numpy(
weights[ptr : ptr + num_b]
).view_as(conv_layer.bias.data)
conv_layer.bias.data.copy_(conv_b)
ptr += num_b
# Load conv. weights
num_w = conv_layer.weight.numel()
conv_w = torch.from_numpy(weights[ptr : ptr + num_w]).view_as(
conv_layer.weight.data
)
conv_layer.weight.data.copy_(conv_w)
ptr += num_w
print("loading weight {}".format(conv_layer))
构建模型以及模型返回值
def _BuildCSPDarknet53(weight_path, resume):
model = CSPDarknet53(weight_path=weight_path, resume=resume)
return model, model.feature_channels[-3:]
到这里,YOLOv4的主干网络就全部定义完成,而这却仅仅只是个开始,后面还有更多的工作需要去完成,目前也只是完成了冰山一角。深度学习,任重而道远,未完待续~
YOLOv4整体网络结构如图:
这里以输入图像尺寸为408*408为例。
该图引用自:睿智的目标检测32——TF2搭建YoloV4目标检测平台(tensorflow2) 不得不说bubbliiing巨佬实在是强
版权声明:本文为CSDN博主「澄鑫」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/ycx_ccc/article/details/122859505
暂无评论