跳至内容

利用YOLOv5中的冻结层进行迁移学习

本指南介绍了如何冻结YOLOv5层的方法。迁移学习是一种强大的机器学习 (ML)技术,它允许您在新数据上快速重新训练模型,而无需从头开始重新训练整个网络。通过冻结初始层的权重,只更新后面层的参数,可以大大减少计算资源需求和训练时间。不过,这种方法可能会略微影响最终模型的准确性

开始之前

首先,克隆YOLOv5 版本库,并安装以下列表中列出的必要依赖项 requirements.txt.确保有 Python.8.0 环境 PyTorch.8 已安装。预先培训 模型 并要求 数据集 将自动从最新的YOLOv5 释放.

git clone https://github.com/ultralytics/yolov5 # clone repository
cd yolov5
pip install -r requirements.txt # install dependencies

冻结层的工作原理

在冷冻层中冷冻时 神经网络在 PyTorch 中,您可以通过设置weights and biases来防止它们的参数weights and biases)在训练过程中被更新。在PyTorch 中,可以通过设置 requires_grad 层的张量属性改为 False.因此,在计算过程中不会对这些层进行梯度计算。 反向传播从而节省计算量和内存。

以下是YOLOv5 如何在训练脚本中实现图层冻结:

# Freeze specified layers
freeze = [f"model.{x}." for x in range(freeze)]  # Define layers to freeze based on module index
for k, v in model.named_parameters():
    v.requires_grad = True  # Ensure all parameters are initially trainable
    if any(x in k for x in freeze):
        print(f"Freezing layer: {k}")
        v.requires_grad = False  # Disable gradient calculation for frozen layers

探索模型架构

了解YOLOv5 模型的结构对于决定冻结哪些层至关重要。您可以使用以下Python 代码段查看所有模块的名称及其参数:

# Assuming 'model' is your loaded YOLOv5 model instance
for name, param in model.named_parameters():
    print(name)

"""
Example Output:
model.0.conv.conv.weight
model.0.conv.bn.weight
model.0.conv.bn.bias
model.1.conv.weight
model.1.bn.weight
model.1.bn.bias
model.2.cv1.conv.weight
model.2.cv1.bn.weight
...
"""

YOLOv5 架构通常包括负责特征提取骨干层(YOLOv5s/m/l/x 等标准配置中的 0-9 层)和执行对象检测的头部(其余各层)。

# Example YOLOv5 v6.0 backbone structure
backbone:
    # [from, number, module, args]
    - [-1, 1, Conv, [64, 6, 2, 2]]  # Layer 0: Initial convolution (P1/2 stride)
    - [-1, 1, Conv, [128, 3, 2]] # Layer 1: Downsampling convolution (P2/4 stride)
    - [-1, 3, C3, [128]]          # Layer 2: C3 module
    - [-1, 1, Conv, [256, 3, 2]] # Layer 3: Downsampling convolution (P3/8 stride)
    - [-1, 6, C3, [256]]          # Layer 4: C3 module
    - [-1, 1, Conv, [512, 3, 2]] # Layer 5: Downsampling convolution (P4/16 stride)
    - [-1, 9, C3, [512]]          # Layer 6: C3 module
    - [-1, 1, Conv, [1024, 3, 2]]# Layer 7: Downsampling convolution (P5/32 stride)
    - [-1, 3, C3, [1024]]         # Layer 8: C3 module
    - [-1, 1, SPPF, [1024, 5]]    # Layer 9: Spatial Pyramid Pooling Fast

# Example YOLOv5 v6.0 head structure
head:
    - [-1, 1, Conv, [512, 1, 1]] # Layer 10
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]] # Layer 11
    - [[-1, 6], 1, Concat, [1]] # Layer 12: Concatenate with backbone P4 (from layer 6)
    - [-1, 3, C3, [512, False]] # Layer 13: C3 module
    # ... subsequent head layers for feature fusion and detection

冷冻选项

您可以使用 --freeze 参数。该参数指定第一个 解冻 模块;该索引之前的所有模块的权重都将被冻结。

仅冻结主干网

冻结整个骨干层(第 0 层到第 9 层),这在调整模型以适应新的对象类别,同时保留从COCO 等大型数据集中学到的一般特征提取功能时很常见:

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 10

当目标数据集与原始训练数据(如 COCO)具有相似的低级视觉特征(边缘、纹理),但包含不同的对象类别时,这种策略就很有效。

冻结除最终检测层以外的所有层

要冻结几乎整个网络,只留下最终输出卷积层(部分 Detect 模块,通常是最后一个模块,如 YOLOv5s 中的模块 24)可训练:

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 24

这种方法适用于主要需要针对不同数量的输出类别调整模型,同时保持绝大多数已学特征不变的情况。它所需的微调计算资源最少。

性能比较

为了说明冷冻层的效果,我们对 YOLOv5m 进行了以下训练 帕斯卡 VOC 数据集 为 50 纪元,从正式的 COCO 预训练开始 重量 (yolov5m.pt).我们比较了三种情况:训练所有层 (--freeze 0),冻结主干 (--freeze 10),冻结除最后探测层以外的所有探测层 (--freeze 24).

# Example command for training with backbone frozen
python train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp hyp.finetune.yaml --freeze 10

准确性结果

结果表明,冻结层可以显著加快训练速度,但可能会导致最终mAP(平均精度)略有下降。对所有层进行训练通常可以获得最佳精度,而冻结更多层可以加快训练速度,但代价可能是降低性能。

比较不同冻结策略的训练 mAP50 结果 训练期间的 mAP50 比较

比较不同冷冻策略的 mAP50-95 训练结果 训练期间 mAP50-95 的比较

性能结果汇总表 绩效指标总表

资源利用

冻结更多层可大幅减少 GPU内存需求和整体利用率。这使得在硬件资源有限的情况下,使用冻结层进行迁移学习成为一种极具吸引力的选择,从而可以训练更大的模型或使用更大的图像尺寸。

训练期间GPU 内存分配百分比 GPU 内存分配 (%)

训练期间GPU 内存使用百分比 GPU 利用率(%)GPU .

何时使用层冷冻

迁移学习过程中的层冻结在几种情况下尤为有利:

  1. 有限的计算资源:如果您的GPU 内存或处理能力有限。
  2. 小数据集:当目标数据集明显小于原始的预训练数据集时,冻结有助于防止过度拟合
  3. 快速原型:当您需要快速调整现有模型,使其适应新任务或新领域,以进行初步评估时。
  4. 相似特征域:如果新数据集中的低层次特征与模型预训练数据集中的特征非常相似,那么就需要对模型进行预训练。

在我们的词汇表条目中探索更多有关迁移学习的细微差别,并考虑采用超参数调整等技术来优化性能。

支持的环境

Ultralytics 提供各种即用型环境,其中包括以下基本依赖项 CUDACuDNNPythonPyTorch已预装。

项目现状

YOLOv5 持续集成状态

此徽章确认所有YOLOv5 GitHub 行动持续集成 (CI) 测试均已成功通过。这些 CI 测试严格评估了YOLOv5 在训练验证推理导出基准等关键操作中的功能和性能。它们确保在 macOS、Windows 和 Ubuntu 上一致、可靠地运行,每 24 小时自动运行一次,并在每次提交新代码时自动运行。



📅创建于 1 年前 ✏️已更新 1 个月前

评论