2023-简单点-yolox-pytorch代码解析（一）-nets/darknet.py

原创已于 2023-12-01 13:33:38 修改 · 1.6w 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #人工智能 #python

于 2023-11-29 10:10:18 首次发布

本文围绕yolox-pytorch展开，给出仓库链接，解析net/darknet.py文件。介绍了BaseConv类，它继承nn.Module类，实现卷积、批归一化和激活函数操作。还阐述了CSPDarknet架构，将输入通道拆分拼接，能有效提取图像特征，增强CNN学习能力，降低计算和内存成本。

部署运行你感兴趣的模型镜像

仓库

https://github.com/bubbliiiing/yolox-pytorch
仓库

yolox网络结构

yolox-pytorch目录

在这里插入图片描述

今天解析注释net/darknet.py

#!/usr/bin/env python3  # 指定使用python3来执行此脚本  
# -*- coding:utf-8 -*-  # 声明脚本使用的编码是utf-8  
# Copyright (c) Megvii, Inc. and its affiliates.  # 版权信息，标注了归属公司为Megvii, Inc.及其附属公司  
  
# 导入torch库，torch是PyTorch的主体库，提供了张量计算和神经网络构建等基础功能  
import torch  
  
# 从torch库中导入nn模块，nn是PyTorch的神经网络模块，提供了各种神经网络层和损失函数等  
from torch import nn  
  ##########################################################################################
# 定义一个名为SiLU的类，继承自nn.Module，这个类代表Sigmoid线性单元，一种非线性激活函数  
class SiLU(nn.Module):  
    # 定义一个静态方法forward，此方法用于定义SiLU类的运算过程，参数x是输入数据  
    @staticmethod  
    def forward(x):   
        # 返回输入数据x乘以sigmoid函数的结果，实现SiLU运算  
        return x * torch.sigmoid(x)  
##########################################################################################  
# 定义一个函数get_activation，用于获取不同类型的激活函数模块  
def get_activation(name="silu", inplace=True):  
    # 根据传入的name参数判断要返回的激活函数类型，默认为SiLU类型  
    if name == "silu":   
        # 如果name为"silu"，则创建一个SiLU类型的模块并赋值给module变量  
        module = SiLU()  
    elif name == "relu":   
        # 如果name为"relu"，则创建一个ReLU类型的模块，inplace参数表示原地操作，是否会修改输入数据，默认为True  
        module = nn.ReLU(inplace=inplace)  
    elif name == "lrelu":   
        # 如果name为"lrelu"，则创建一个LeakyReLU类型的模块，0.1表示负斜率，inplace参数表示原地操作，是否会修改输入数据，默认为True  
        module = nn.LeakyReLU(0.1, inplace=inplace)  
    else:   
        # 如果name不是上述任何一种类型，则抛出一个AttributeError异常，提示"Unsupported act type: {}".format(name)错误信息  
        raise AttributeError("Unsupported act type: {}".format(name))  
    # 返回获取到的激活函数模块  
    return module  
##########################################################################################  
# 定义一个名为Focus的类，继承自nn.Module，这个类代表焦点模块，用于对输入数据进行空间上的重新构造  
class Focus(nn.Module):  
    # 定义一个构造函数__init__，此方法用于初始化Focus类的实例对象  
    def __init__(self, in_channels, out_channels, ksize=1, stride=1, act="silu"):  
        # 调用父类的构造函数，进行基本的初始化操作  
        super().__init__()  
        # 定义变量pad为(ksize - 1) // 2，表示卷积操作的填充大小  
        pad = (ksize - 1) // 2  
        # 创建一个BaseConv类型的模块，参数为in_channels * 4（输入通道数变为原来的四倍），out_channels（输出通道数），ksize（卷积核大小），stride（步长），pad（填充大小），act（激活函数类型默认为SiLU）  
        self.conv = BaseConv(in_channels * 4, out_channels, ksize, stride, act=act)  
    # 定义一个forward方法，此方法用于定义Focus类的运算过程，参数x是输入数据  
    def forward(self, x):   
        # 从输入数据x中获取四个位置的patch并进行拼接，形成一个新的数据x并返回给conv进行卷积操作
        #考虑2 x 2的四个方格，左上，左下，右上和右下为起始点，各一个元素采样  
        patch_top_left = x[..., ::2, ::2]  # 从左上角开始每隔一个像素取一个像素点形成左上角patch  
        patch_bot_left = x[..., 1::2, ::2]  # 从左下角开始每隔一个像素取一个像素点形成左下角patch  
        patch_top_right = x[..., ::2, 1::2]  # 从右上角开始每隔一个像素取一个像素点形成右上角patch  
        patch_bot_right = x[..., 1::2, 1::2]  # 从右下角开始每隔一个像素取一个像素点形成右下角patch  
        x = torch.cat((patch_top_left, patch_bot_left, patch_top_right,patch_bot_right), dim=1)  # 将四个patch在通道维度上进行拼接  
        # 返回经过conv卷积操作后的结果  
        return self.conv(x)  
##########################################################################################  
# 定义一个名为BaseConv的类，继承自nn.Module，这个类代表基础卷积模块，用于构建卷积神经网络的基础模块  
class BaseConv(nn.Module):  
    # 定义一个构造函数__init__，此方法用于初始化BaseConv类的实例对象  
    def __init__(self, in_channels, out_channels, ksize, stride, groups=1, bias=False, act="silu"):  
        # 调用父类的构造函数，进行基本的初始化操作  
        super().__init__()  
        # 定义变量pad为(ksize - 1) // 2，表示卷积操作的填充大小  
        pad = (ksize - 1) // 2  
        # 创建一个Conv2d类型的卷积层，参数依次为输入通道数，输出通道数，卷积核大小，步长，填充大小，分组数，是否使用偏置  
        self.conv = nn.Conv2d(in_channels, out_channels, ksize, stride, pad, groups=groups, bias=bias)
        # 定义一个批归一化层，用于加速训练和提高模型稳定性。  
        self.bn   = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.03)   
        # 如果传入的act参数不为空字符串，则根据act参数创建一个激活函数模块并赋值给self.act变量  
        if act is not None:  
            self.act = get_activation(act, inplace=True)  
        else:  
            self.act = nn.Identity()  # 如果act参数为空字符串，则使用恒等映射作为激活函数  
    # 定义一个forward方法，此方法用于定义BaseConv类的运算过程，参数x是输入数据  
    def forward(self, x):   
        # 返回经过conv卷积操作和act激活函数处理后的结果  
        return self.act(self.conv(x))
#########################################################################################

# 定义一个名为DWConv的类，继承自nn.Module，这是一个深度可分离卷积模块。  
class DWConv(nn.Module):  
    # 初始化函数，用于设置该模块所需的各种参数和子模块。  
    def __init__(self, in_channels, out_channels, ksize, stride=1, act="silu"):  
        # 调用父类的初始化函数。  
        super().__init__()  
        # 定义一个深度卷积层（逐通道卷积），其输出通道数与输入通道数相同，并且使用给定的核大小和步长。  
        # 注意这里的groups参数设置为in_channels，意味着每个输入通道都有独立的卷积核。  
        self.dconv = BaseConv(in_channels, in_channels, ksize=ksize, stride=stride, groups=in_channels, act=act,)  
        # 定义一个逐点卷积层（1x1卷积），用于改变通道数。这里的groups参数设置为1，表示是普通的卷积操作。  
        self.pconv = BaseConv(in_channels, out_channels, ksize=1, stride=1, groups=1, act=act)  
  
    # 定义前向传播函数，输入数据x会首先经过深度卷积层处理，然后再经过逐点卷积层处理。  
    def forward(self, x):  
        x = self.dconv(x)  # 数据经过深度卷积层处理。  
        return self.pconv(x)  # 经过逐点卷积层处理后输出结果。
##########################################################################################
# 定义一个名为SPPBottleneck的类，继承自nn.Module，这是一个包含空间金字塔池化（Spatial Pyramid Pooling, SPP）的瓶颈模块。  
class SPPBottleneck(nn.Module):  
    # 初始化函数，用于设置该模块所需的各种参数和子模块。  
    def __init__(self, in_channels, out_channels, kernel_sizes=(5, 9, 13), activation="silu"):  
        # 调用父类的初始化函数。  
        super().__init__()  
        # 计算隐藏层的通道数，为输入通道数的一半。  
        hidden_channels = in_channels // 2  
        # 定义第一个卷积层，用于降低通道数。这里使用1x1的卷积核。  
        self.conv1      = BaseConv(in_channels, hidden_channels, 1, stride=1, act=activation)  
        # 定义一个包含多个最大池化层的模块列表，用于生成不同尺度的空间金字塔特征。  
        self.m          = nn.ModuleList([nn.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2) for ks in kernel_sizes])  
        # 计算第二个卷积层的输入通道数，为隐藏层通道数与空间金字塔层数的和。  
        conv2_channels  = hidden_channels * (len(kernel_sizes) + 1)  
        # 定义第二个卷积层，用于增加通道数并整合空间金字塔特征。这里使用1x1的卷积核。  
        self.conv2      = BaseConv(conv2_channels, out_channels, 1, stride=1, act=activation)  
  
    # 定义前向传播函数，输入数据x会首先经过第一个卷积层处理，然后生成空间金字塔特征并拼接在一起，最后经过第二个卷积层处理并输出结果。  
    def forward(self, x):  
        x = self.conv1(x)  # 数据经过第一个卷积层处理。  
        x = torch.cat([x] + [m(x) for m in self.m], dim=1)  # 生成空间金字塔特征并将它们拼接在一起。  
        x = self.conv2(x)  # 经过第二个卷积层处理后输出结果。  
        return x
##########################################################################################
class Bottleneck(nn.Module):  # 定义一个名为Bottleneck的类，它继承了nn.Module  
    # Standard bottleneck  
    def __init__(self, in_channels, out_channels, shortcut=True, expansion=0.5, depthwise=False, act="silu",):  # 初始化函数  
        super().__init__()  # 调用父类的初始化函数  
          
        # 计算隐藏层的通道数，通过输出通道数与扩展因子相乘并取整得到  
        hidden_channels = int(out_channels * expansion)    
          
        # 根据depthwise的值选择卷积类型，如果是True则使用深度可分离卷积(DWConv)，否则使用基础卷积(BaseConv)  
        Conv = DWConv if depthwise else BaseConv  
          
        # 利用1x1卷积进行通道数的缩减，通常缩减率是50%  
        self.conv1 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)  
          
        # 利用3x3卷积进行通道数的扩展，并完成特征提取  
        self.conv2 = Conv(hidden_channels, out_channels, 3, stride=1, act=act)  
          
        # 判断是否使用残差连接，只有当shortcut为True且输入和输出通道数相等时才使用  
        self.use_add = shortcut and in_channels == out_channels  
  
    def forward(self, x):  # 定义前向传播函数  
        # 首先应用第一个卷积层（通道数缩减），然后应用第二个卷积层（通道数扩展）  
        y = self.conv2(self.conv1(x))  
          
        # 如果满足使用残差连接的条件，则将输入x添加到y中  
        if self.use_add:  
            y = y + x  
              
        # 返回y作为模块的输出  
        return y

##########################################################################################
class CSPLayer(nn.Module):  
    def __init__(self, in_channels, out_channels, n=1, shortcut=True, expansion=0.5, depthwise=False, act="silu",):  
        super().__init__()  
        hidden_channels = int(out_channels * expansion)  
  
        # 主干部分的初次卷积，用于通道数的缩减和特征初步提取  
        self.conv1 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)  
  
        # 大的残差边部分的初次卷积，与conv1有相同的通道数和作用，但其输出会与后续的Bottleneck结构进行堆叠  
        self.conv2 = BaseConv(in_channels, hidden_channels, 1, stride=1, act=act)  
  
        # 对堆叠的结果进行卷积的处理，将2倍的隐藏通道数转换为输出通道数  
        self.conv3 = BaseConv(2 * hidden_channels, out_channels, 1, stride=1, act=act)  
  
        # 根据循环的次数构建上述Bottleneck残差结构，多次重复Bottleneck结构可以增强特征提取能力  
        module_list = [Bottleneck(hidden_channels, hidden_channels, shortcut, 1.0, depthwise, act=act) for _ in range(n)]  
        self.m = nn.Sequential(*module_list)  # 使用Sequential将多个Bottleneck结构串联起来
	
	def forward(self, x):
	        #-------------------------------#
	        #   x_1是主干部分
	        #-------------------------------#
	        x_1 = self.conv1(x)
	        #-------------------------------#
	        #   x_2是大的残差边部分
	        #-------------------------------#
	        x_2 = self.conv2(x)
	
	        #-----------------------------------------------#
	        #   主干部分利用残差结构堆叠继续进行特征提取
	        #-----------------------------------------------#
	        x_1 = self.m(x_1)
	        #-----------------------------------------------#
	        #   主干部分和大的残差边部分进行堆叠
	        #-----------------------------------------------#
	        x = torch.cat((x_1, x_2), dim=1)
	        #-----------------------------------------------#
	        #   对堆叠的结果进行卷积的处理
	        #-----------------------------------------------#
	        return self.conv3(x)
##########################################################################################  

class CSPDarknet(nn.Module):  
    def __init__(self, dep_mul, wid_mul, out_features=("dark3", "dark4", "dark5"), depthwise=False, act="silu",):  
        super().__init__()  
        assert out_features, "please provide output features of Darknet"  # 确保提供了Darknet的输出特征  
        self.out_features = out_features  
        Conv = DWConv if depthwise else BaseConv  # 根据depthwise的值选择卷积类型  
  
        # 输入图片的基本信息，初始的基本通道数是64，乘上宽度因子wid_mul得到实际的初始通道数  
        base_channels = int(wid_mul * 64)  # 64  
        base_depth = max(round(dep_mul * 3), 1)  # 深度是3乘以深度因子dep_mul，至少为1  
  
        # 利用focus网络结构进行特征提取，将输入图片从640x640x3变换到320x320x64  
        self.stem = Focus(3, base_channels, ksize=3, act=act)  
  
        # dark2部分 -----------------------------------------------------  
        # 完成卷积之后，特征图尺寸从320x320x64变为160x160x128  
        # 完成CSPlayer之后，特征图尺寸保持为160x160x128不变  
        self.dark2 = nn.Sequential(  
            Conv(base_channels, base_channels * 2, 3, 2, act=act),  # 卷积操作，通道数翻倍，高度和宽度减半  
            CSPLayer(base_channels * 2, base_channels * 2, n=base_depth, depthwise=depthwise, act=act),  # CSPLayer结构  
        )  
  
        # dark3部分 -----------------------------------------------------  
        # 完成卷积之后，特征图尺寸从160x160x128变为80x80x256  
        # 完成CSPlayer之后，特征图尺寸保持为80x80x256不变  
        self.dark3 = nn.Sequential(  
            Conv(base_channels * 2, base_channels * 4, 3, 2, act=act),  # 卷积操作，通道数翻倍，高度和宽度减半  
            CSPLayer(base_channels * 4, base_channels * 4, n=base_depth * 3, depthwise=depthwise, act=act),  # CSPLayer结构  
        )  
        # 完成卷积之后，80x80x256的特征图变为40x40x512  
		# 完成CSPlayer之后，特征图的尺寸和通道数保持不变，为40x40x512  
		self.dark4 = nn.Sequential(  
		    Conv(base_channels * 4, base_channels * 8, 3, 2, act=act),  # 卷积操作，通道数翻倍，高度和宽度减半  
		    CSPLayer(base_channels * 8, base_channels * 8, n=base_depth * 3, depthwise=depthwise, act=act),  # CSPLayer结构  
		)  
		  
		# 完成卷积之后，40x40x512的特征图变为20x20x1024  
		# 完成SPP（Spatial Pyramid Pooling）模块之后，特征图的尺寸和通道数保持不变，为20x20x1024  
		# 完成CSPlayer之后，特征图的尺寸和通道数仍然保持不变，为20x20x1024  
		self.dark5 = nn.Sequential(  
		    Conv(base_channels * 8, base_channels * 16, 3, 2, act=act),  # 卷积操作，通道数翻倍，高度和宽度减半  
		    SPPBottleneck(base_channels * 16, base_channels * 16, activation=act),  # SPPBottleneck结构，增强特征的多样性  
		    CSPLayer(base_channels * 16, base_channels * 16, n=base_depth, shortcut=False, depthwise=depthwise, act=act),  # CSPLayer结构  
		)
	def forward(self, x):  
	    outputs = {}  # 初始化一个空字典，用于存储每一层的输出  
	  
	    x = self.stem(x)  # stem层，通常是模型的初始部分，用于处理输入数据  
	    outputs["stem"] = x  # 将stem层的输出存储到字典中  
	  
	    x = self.dark2(x)  # dark2层，是模型的一部分，用于提取特征  
	    outputs["dark2"] = x  # 将dark2层的输出存储到字典中  
	  
	    # dark3的输出为80, 80, 256，是一个有效特征层  
	    x = self.dark3(x)  # dark3层，继续提取特征  
	    outputs["dark3"] = x  # 将dark3层的输出存储到字典中  
	  
	    # dark4的输出为40, 40, 512，是一个有效特征层  
	    x = self.dark4(x)  # dark4层，继续提取特征  
	    outputs["dark4"] = x  # 将dark4层的输出存储到字典中  
	  
	    # dark5的输出为20, 20, 1024，是一个有效特征层  
	    x = self.dark5(x)  # dark5层，继续提取特征  
	    outputs["dark5"] = x  # 将dark5层的输出存储到字典中  
	  
	     
	    # 只返回out_features列表中指定的层的输出；
	    #在代码的最后一行，使用了字典推导式来过滤outputs字典中的项，
	    #只保留self.out_features中指定的键对应的值。
	    #这意味着最终返回的将是out_features列表中指定的层的输出。
	    #这是一种常见的做法，可以让用户灵活选择需要输出的特征层。
	    return {k: v for k, v in outputs.items() if k in self.out_features}

Focus

在这里插入图片描述

BaseConv

简而言之：卷积 + 激活【是否】

BaseConv是一个基础卷积类，它继承了PyTorch中的nn.Module类，并实现了卷积、批归一化（Batch Normalization）和激活函数等核心操作。BaseConv类的主要参数包括：

in_channels：输入通道数。
out_channels：输出通道数。
ksize：卷积核大小。
stride：步长。
groups：分组卷积中的组数，默认为1。
bias：是否使用bias（偏差），默认为False。
act：激活函数类型，默认为"silu"。

在BaseConv类中，主要实现了以下三个方法：

__init__：构造函数，用于初始化BaseConv对象。在构造函数中，会创建一个nn.Conv2d对象（卷积层），一个nn.BatchNorm2d对象（批归一化层）和一个激活函数对象。
forward：前向传播函数。输入数据x首先经过卷积层和批归一化层，然后通过激活函数进行激活，最终输出结果。
get_activation：获取激活函数。该函数用于获取指定名称的激活函数对象。

总体来说，BaseConv是一个简单但实用的卷积类，可以作为构建其他复杂卷积网络的基础组件。

DWConv

在这里插入图片描述

SPPBottleneck

在这里插入图片描述

CSPDarknet

CSP:Cross Stage Partial
在这里插入图片描述
CSP darknet的架构具有以下特点：

它将输入通道分成两部分，一部分经过常规Res(X)Block后跟另一部分进行通道拼接，拼接后进入transition layer。CSP将梯度的变化从头到尾地集成到特征图中，在减少了计算量的同时可以保证准确率（或有所提升）。其中transition layer是一个卷积加池化，用于整合学到的特征，降低特征图的尺寸。
CSP模块中的每个模块都有特定的特点，例如输出feature map大小减半，通道数增倍等。

CSP darknet的这种架构可以有效地提取输入图像的特征，并为后续的网络处理提供有效的特征集合。

优势