YOLOv4代码详解笔记-Backbone
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巨佬实在是强