YOLOv6-4.0部分代码阅读笔记-yolo_lite.py
yolo_lite.py
yolov6\models\yolo_lite.py
所需的库和模块
#!/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.reppan import *
from yolov6.models.efficientrep import *
from yolov6.utils.events import LOGGER
from yolov6.models.heads.effidehead_lite import Detect, build_effidehead_layerclass 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 ：目标检测任务中的类别数。def __init__(self, config, channels=3, num_classes=None): # model, input channels, number of classes 模型、输入通道、类别数量# 调用父类初始化方法，调用 nn.Module 的初始化方法，这是 PyTorch 中定义模型类时的标准做法。super().__init__()# Build network 建立网络# def build_network(config, in_channels, num_classes): -> 函数返回构建好的 骨干网络 、 颈部网络 和 头部网络 。 -> return backbone, neck, headself.backbone, self.neck, self.detect = build_network(config, channels, num_classes)# Init Detect head 初始检测头# 从头部网络（detect）获取模型的步长（stride），步长是目标检测中的一个重要参数，它决定了检测的精度和召回率。self.stride = self.detect.stride# def initialize_biases(self): -> 它用于初始化神经网络中特定层的偏置（biases）。这个方法特别针对于类别预测层（ self.cls_preds ）和边界框回归预测层（ self.reg_preds ）的偏置和权重进行初始化。# 对检测头的偏置进行初始化，这通常是根据特定的初始化策略来设置偏置值，以帮助模型在训练初期更快地收敛。self.detect.initialize_biases()# Init weights 初始权重# def initialize_weights(model): -> 它用于初始化传入的模型（ model ）中的权重和偏置。这个函数遍历模型中的所有模块，并根据模块的类型应用特定的初始化策略。# 对整个模型的权重进行初始化，这通常包括对卷积层、全连接层等的权重进行初始化，以确保模型在训练开始时有一个良好的起点。initialize_weights(self)def forward(self, x):# 检查当前是否处于将 PyTorch 模型导出为 ONNX 格式的过程。# ONNX（Open Neural Network Exchange）是一种开放的格式，用于表示深度学习模型，以便在不同的框架和工具之间进行模型的交换和部署。export_mode = torch.onnx.is_in_onnx_export()# 输入数据 x 首先通过骨干网络（backbone）进行处理。x = self.backbone(x)# 骨干网络的输出接着被送入颈部网络（neck）进行进一步的处理。x = self.neck(x)# 如果当前不在 ONNX 导出模式，那么代码会执行以下操作。if export_mode == False:# 初始化一个空列表 featmaps ，用于存储特征图。featmaps = []# 将颈部网络的输出（ x ）添加到 featmaps 列表中。这里的 x 可能包含多个特征图，这些特征图可以用于后续的目标检测或其他任务。featmaps.extend(x)# 颈部网络的输出被送入检测头（detect），进行目标检测。x = self.detect(x)# 根据是否处于 ONNX 导出模式返回不同的结果。# 如果处于 ONNX 导出模式（ export_mode is True ），则只返回检测头的输出 x 。# 如果不在 ONNX 导出模式，返回一个包含检测头输出 x 和特征图列表 featmaps 的列表。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 selfdef build_network(config, in_channels, num_classes):
def build_network(config, in_channels, num_classes):width_mul = config.model.width_multiple # 控制网络结构宽度的缩放因子。num_repeat_backbone = config.model.backbone.num_repeats # 主干网络每个stage中基础模块的重复个数。out_channels_backbone = config.model.backbone.out_channels # 主干网络每个stage中输出的通道数。# scale_size_backbone ：骨干网络尺度缩放因子。scale_size_backbone = config.model.backbone.scale_size# in_channels_neck ：颈部网络输入通道数的列表。in_channels_neck = config.model.neck.in_channels# unified_channels_neck ：颈部网络统一通道数。unified_channels_neck = config.model.neck.unified_channelsin_channels_head = config.model.head.in_channels # 检测头每个特征层的输入通道数。num_layers = config.model.head.num_layers # 检测头的特征层数量, P6模型此值为4。BACKBONE = eval(config.model.backbone.type) # 主干网络的类别，目前可支持'EfficientRep', 'CSPBepBackbone','EfficientRep6','CSPBepBackbone_P6' 4种。NECK = eval(config.model.neck.type) # 检测器 Neck 的类别，目前可选用'RepPANNeck', 'CSPRepPANNeck','RepBiFPANNeck','CSPRepBiFPANNeck','RepBiFPANNeck6','CSPRepBiFPANNeck_P6' 6种。# def make_divisible(v, divisor=16): -> return new_v# 使用 make_divisible 函数调整骨干网络和颈部网络的通道数，使其可以被特定的除数整除，这里通常是 16 或 8，以优化模型的性能和参数数量。# out_channels_backbone ：骨干网络输出通道数调整后的结果。out_channels_backbone = [make_divisible(i * width_mul)for i in out_channels_backbone]# mid_channels_backbone ：骨干网络中间通道数，根据 scale_size_backbone 缩放后的结果。mid_channels_backbone = [make_divisible(int(i * scale_size_backbone), divisor=8) for i in out_channels_backbone]# in_channels_neck ：颈部网络输入通道数调整后的结果。in_channels_neck = [make_divisible(i * width_mul)for i in in_channels_neck]# 实例化骨干网络，传入输入 通道数 、 中间通道数 、 输出通道数 和 重复次数 。backbone = BACKBONE(in_channels,mid_channels_backbone,out_channels_backbone,num_repeat=num_repeat_backbone)# 实例化颈部网络，传入颈部网络 输入通道数 和 统一通道数 。neck = NECK(in_channels_neck, unified_channels_neck)# def build_effidehead_layer(channels_list, num_anchors, num_classes, num_layers): -> 用于构建头部网络的层结构。 -> return head_layershead_layers = build_effidehead_layer(in_channels_head, 1, num_classes, num_layers)# class Detect(nn.Module): -> def __init__(self, num_classes=80, num_layers=3, inplace=True, head_layers=None):# head ：实例化头部网络 ， 传入类别数 、 层数 和 头部网络的层结构 。head = Detect(num_classes, num_layers, head_layers=head_layers)# 函数返回构建好的 骨干网络 、 颈部网络 和 头部网络 。return backbone, neck, headdef build_model(cfg, num_classes, device):
def build_model(cfg, num_classes, device):model = Model(cfg, channels=3, num_classes=num_classes).to(device)return modeldef make_divisible(v, divisor=16):
# 这个函数的作用是将一个给定的数值 v 调整为最接近的、可以被 divisor 整除的数值。在深度学习模型中，这样的函数通常用于确保模型的某些参数（比如卷积层的输出通道数）可以被某个特定的数整除，以便于优化内存使用和计算效率。
def make_divisible(v, divisor=16):# 首先， v 是需要调整的原始数值。# divisor / 2 是为了避免在四舍五入时产生误差，所以先加上 divisor 的一半。# int(v + divisor / 2) 将加上 divisor 一半后的值转换为整数。# // divisor 是整数除法，确保结果可以被 divisor 整除。# * divisor 将结果乘以 divisor ，以确保数值是 divisor 的整数倍。# max(divisor, ...) 确保 new_v 至少为 divisor ，避免结果太小。new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)# 这个条件检查调整后的 new_v 是否小于原始 v 的 90%。 如果是，说明调整后的数值相对于原始数值减少得太多，可能对模型性能有影响。if new_v < 0.9 * v:# 如果 new_v 小于 0.9 * v ，则将 divisor 加到 new_v 上，以确保调整后的数值不会比原始数值小太多。new_v += divisor# 函数返回调整后的 new_v 。return new_v