yolo.py
yolov6\models\yolo.py
目录
yolo.py
1.所需的库和模块
2.class Model(nn.Module):
3.def make_divisible(x, divisor):
4.def build_network(config, channels, num_classes, num_layers, fuse_ab=False, distill_ns=False):
5.def build_model(cfg, num_classes, device, fuse_ab=False, distill_ns=False):
1.所需的库和模块
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from yolov6.layers.common import *
from yolov6.utils.torch_utils import initialize_weights
from yolov6.models.efficientrep import *
from yolov6.models.reppan import *
from yolov6.utils.events import LOGGER
2.class Model(nn.Module):
class Model(nn.Module):# 具有主干、颈部和头部的 YOLOv6 模型。# 默认部件是 EfficientRep Backbone ， Rep-PAN 和 Efficient Decoupled Head 。'''YOLOv6 model with backbone, neck and head.The default parts are EfficientRep Backbone, Rep-PAN andEfficient Decoupled Head.'''# 1.config ：模型配置对象，包含了模型结构和训练相关的参数。# 2.channels ：输入图像的通道数，默认为3（彩色图像）。# 3.num_classes ：目标检测任务中的类别数。# 4.fuse_ab ：是否融合 A 和 B 特征层。# 5.distill_ns ：是否进行蒸馏。def __init__(self, config, channels=3, num_classes=None, fuse_ab=False, distill_ns=False): # model, input channels, number of classes# 调用父类的初始化方法，调用 nn.Module 的初始化方法，这是 PyTorch 中定义模型类时的标准做法。super().__init__()# Build network 建立网络# 从配置中获取头部网络的层数。num_layers = config.model.head.num_layers# 调用 build_network 函数构建模型的三个部分，并将返回的 骨干网络 、 颈部网络 和 头部网络 分别赋值给 self.backbone 、 self.neck 和 self.detect 。self.backbone, self.neck, self.detect = build_network(config, channels, num_classes, num_layers, fuse_ab=fuse_ab, distill_ns=distill_ns)# Init Detect head 初始检测头# 从头部网络（detect）获取模型的步长（stride），步长是目标检测中的一个重要参数，它决定了检测的精度和召回率。self.stride = self.detect.stride# 对检测头的偏置进行初始化，这通常是根据特定的初始化策略来设置偏置值，以帮助模型在训练初期更快地收敛。self.detect.initialize_biases()# Init weights 初始权重# 对整个模型的权重进行初始化，这通常包括对卷积层、全连接层等的权重进行初始化，以确保模型在训练开始时有一个良好的起点。initialize_weights(self)def forward(self, x):export_mode = torch.onnx.is_in_onnx_export()x = self.backbone(x)x = self.neck(x)if export_mode == False:featmaps = []featmaps.extend(x)x = self.detect(x)return x if export_mode is True else [x, featmaps]# 这段代码定义了 Model 类的 _apply 方法，这个方法在 PyTorch 中用于对模型中的参数和缓冲区应用一个函数 fn 。这通常用于模型的变换，比如参数的修改或者参数的复制等。def _apply(self, fn):# 首先调用 nn.Module 的 _apply 方法，这个方法会对模型中的所有参数和缓存应用函数 fn 。这里的 self 被重新赋值为 _apply 方法的返回值，这意味着模型的参数已经被 fn 函数处理过了。self = super()._apply(fn)# 对检测头中的 stride 属性应用函数 fn 。 stride 是一个重要的参数，它决定了模型输出的空间分辨率。通过 fn 函数，可以对 stride 进行调整或变换。self.detect.stride = fn(self.detect.stride)# 对检测头中的 grid 属性应用函数 fn 。 grid 属性通常是一个列表或者张量，包含了在不同尺度上的目标位置信息。通过 map 函数， fn 被应用到 grid 的每一个元素上，然后返回一个新的列表。self.detect.grid = list(map(fn, self.detect.grid))# 返回经过参数变换后的模型实例。return self3.def make_divisible(x, divisor):
# 将给定的数值 x 调整为最接近的、可以被 divisor 整除的数值。
# 1.x ：需要调整的原始数值。
# 2.divisor ：用于调整 x 的除数，即 x 需要被这个数整除。
def make_divisible(x, divisor):# 向上修正 x 的值，使其可以被除数整除。# Upward revision the value x to make it evenly divisible by the divisor.# math.ceil(x)# 方法将 x 向上舍入到最接近的整数。# x / divisor ：首先计算 x 除以 divisor 的结果，得到一个浮点数。# math.ceil(x / divisor) ：使用 math.ceil 函数对上述结果进行向上取整，即取大于或等于该浮点数的最小整数。# math.ceil(x / divisor) * divisor ：将向上取整后的结果乘以 divisor ，得到一个可以被 divisor 整除的数值。return math.ceil(x / divisor) * divisor4.def build_network(config, channels, num_classes, num_layers, fuse_ab=False, distill_ns=False):
def build_network(config, channels, num_classes, num_layers, fuse_ab=False, distill_ns=False):depth_mul = config.model.depth_multiple # 控制网络结构深度的缩放因子width_mul = config.model.width_multiple # 控制网络结构宽度的缩放因子num_repeat_backbone = config.model.backbone.num_repeats # 主干网络每个stage中基础模块的重复个数channels_list_backbone = config.model.backbone.out_channels # 主干网络每个stage中输出的通道数fuse_P2 = config.model.backbone.get('fuse_P2') # 是否融合骨干网络P2层特征cspsppf = config.model.backbone.get('cspsppf') # 是否使用CSPSPPF模块以替换SPPF模块num_repeat_neck = config.model.neck.num_repeats # Neck网络连接每个特征层基础模块的重复个数channels_list_neck = config.model.neck.out_channels # Neck网络连接每个特征层的上采样/下采样模块通道数use_dfl = config.model.head.use_dfl # 是否使用 distributed focal loss ，若后续想继续蒸馏，需要设为True以保留DFL分支reg_max = config.model.head.reg_max # 若 use_dfl 为 False, 则 reg_max 设为 0；若 use_dfl 为 True, 则reg_max 设为 16num_repeat = [(max(round(i * depth_mul), 1) if i > 1 else i) for i in (num_repeat_backbone + num_repeat_neck)]# def make_divisible(x, divisor): -> 向上修正 x 的值，使其可以被除数整除。 -> return math.ceil(x / divisor) * divisor# 使用 make_divisible 函数调整骨干网络和颈部网络的通道数，使其可以被 8 整除。channels_list = [make_divisible(i * width_mul, 8) for i in (channels_list_backbone + channels_list_neck)]# 根据训练模式获取相应的网络模块。block = get_block(config.training_mode)# 根据配置动态加载骨干网络的类。BACKBONE = eval(config.model.backbone.type) # 网络类型# 根据配置动态加载颈部网络的类。NECK = eval(config.model.neck.type) # 检测器 Neck 的类别，目前可选用'RepPANNeck', 'CSPRepPANNeck','RepBiFPANNeck','CSPRepBiFPANNeck','RepBiFPANNeck6','CSPRepBiFPANNeck_P6' 6种if 'CSP' in config.model.backbone.type: # 主干网络的类别，目前可支持'EfficientRep', 'CSPBepBackbone','EfficientRep6','CSPBepBackbone_P6' 4种if "stage_block_type" in config.model.backbone:stage_block_type = config.model.backbone.stage_block_typeelse:stage_block_type = "BepC3" #default 默认backbone = BACKBONE(in_channels=channels,channels_list=channels_list,num_repeats=num_repeat,block=block,csp_e=config.model.backbone.csp_e,fuse_P2=fuse_P2,cspsppf=cspsppf,stage_block_type=stage_block_type)neck = NECK(channels_list=channels_list,num_repeats=num_repeat,block=block,csp_e=config.model.neck.csp_e,stage_block_type=stage_block_type)else:backbone = BACKBONE(in_channels=channels,channels_list=channels_list,num_repeats=num_repeat,block=block,fuse_P2=fuse_P2,cspsppf=cspsppf)neck = NECK(channels_list=channels_list,num_repeats=num_repeat,block=block)# 它处理了在模型中启用蒸馏（distillation）时头部网络的构建。if distill_ns:# 如果 distill_ns 参数为 True ，则从 yolov6.models.heads.effidehead_distill_ns 模块导入 Detect 类和 build_effidehead_layer 函数。这些是专门用于蒸馏的头部网络组件。from yolov6.models.heads.effidehead_distill_ns import Detect, build_effidehead_layerif num_layers != 3:# 如果 num_layers 不等于 3，使用 LOGGER.error 记录错误信息，并退出程序。这是因为在蒸馏模式下，模型的头部网络层数需要与特定的配置相匹配，这里假设蒸馏模型的头部网络层数固定为 3。# 错误：蒸馏模式不适合带有 P6 头的 n/s 型号。LOGGER.error('ERROR in: Distill mode not fit on n/s models with P6 head.\n')# exit()# 结束整个程序。# 在python中运行一段代码，如果在某处已经完成整次任务，可以用 exit 退出整个运行。并且还可以在 exit() 的括号里加入自己退出程序打印说明。exit()# def build_effidehead_layer(channels_list, num_anchors, num_classes, reg_max=16): -> return head_layers# 调用 build_effidehead_layer 函数来构建头部网络的层结构。它根据输入的通道数列表 channels_list 、锚框数量（这里为 1）、类别数 num_classes 和 reg_max 来构建头部网络层。head_layers = build_effidehead_layer(channels_list, 1, num_classes, reg_max=reg_max)# class Detect(nn.Module): -> 高效分离头，实现免费蒸馏。（适用于纳米/小型型号）。 -> def __init__(self, num_classes=80, num_layers=3, inplace=True, head_layers=None, use_dfl=True, reg_max=16): # detection layer# 使用 Detect 类创建头部网络实例。这里传入类别数 num_classes 、头部网络层数 num_layers 、头部网络层结构 head_layers 和是否使用分布式焦点损失的标志 use_dfl 。head = Detect(num_classes, num_layers, head_layers=head_layers, use_dfl=use_dfl)# 这段代码处理了在模型中启用特征融合（fuse A and B features）时头部网络的构建。elif fuse_ab:# 如果 fuse_ab 参数为 True ，则从 yolov6.models.heads.effidehead_fuseab 模块导入 Detect 类和 build_effidehead_layer 函数。这些是专门用于融合特征的头部网络组件。from yolov6.models.heads.effidehead_fuseab import Detect, build_effidehead_layer# 从配置中获取初始化锚点（anchors），这些锚点用于目标检测中的目标框定位。anchors_init = config.model.head.anchors_inithead_layers = build_effidehead_layer(channels_list, 3, num_classes, reg_max=reg_max, num_layers=num_layers)# class Detect(nn.Module): -> 高效的分离头，用于融合锚框分支。 -> def __init__(self, num_classes=80, anchors=None, num_layers=3, inplace=True, head_layers=None, use_dfl=True, reg_max=16): # detection layerhead = Detect(num_classes, anchors_init, num_layers, head_layers=head_layers, use_dfl=use_dfl)else:from yolov6.models.effidehead import Detect, build_effidehead_layerhead_layers = build_effidehead_layer(channels_list, 1, num_classes, reg_max=reg_max, num_layers=num_layers)# class Detect(nn.Module): -> 高效分离头。利用硬件感知设计，使用混合通道方法对解耦头进行优化。 -> def __init__(self, num_classes=80, num_layers=3, inplace=True, head_layers=None, use_dfl=True, reg_max=16): # detection layer 检测层head = Detect(num_classes, num_layers, head_layers=head_layers, use_dfl=use_dfl)return backbone, neck, head5.def build_model(cfg, num_classes, device, fuse_ab=False, distill_ns=False):
def build_model(cfg, num_classes, device, fuse_ab=False, distill_ns=False):# class Model(nn.Module): -> 具有主干、颈部和头部的 YOLOv6 模型。默认部件是 EfficientRep Backbone ， Rep-PAN 和 Efficient Decoupled Head 。# -> def __init__(self, config, channels=3, num_classes=None, fuse_ab=False, distill_ns=False): # model, input channels, number of classesmodel = Model(cfg, channels=3, num_classes=num_classes, fuse_ab=fuse_ab, distill_ns=distill_ns).to(device)return model
![[论文阅读]LOGAN: Membership Inference Attacks Against Generative Models](http://pic.xiahunao.cn/[论文阅读]LOGAN: Membership Inference Attacks Against Generative Models)
