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整体网络结构如图：

该图引用自：深入浅出Yolo系列之Yolov3&Yolov4&Yolov5&Yolox核心基础知识完整讲解

这里以输入图像尺寸为408*408为例。

该图引用自：睿智的目标检测32——TF2搭建YoloV4目标检测平台（tensorflow2） 不得不说bubbliiing巨佬实在是强