YOLO Vision 2026:

YOLOv5 中使用冻结层进行迁移学习#

📚 本指南介绍如何在实现迁移学习冻结 YOLOv5 🚀 层。迁移学习是一种强大的机器学习 (ML)技术,可以让你在新数据上快速重新训练模型,而无需从头开始重新训练整个网络。通过冻结初始层的权重,仅更新后续层的参数,你可以显著降低计算资源需求和训练时间。不过,这种方法可能会对最终模型的准确率造成轻微影响。

开始之前#

首先,克隆 YOLOv5 仓库,并安装 requirements.txt 中列出的必要依赖项。确保你使用的环境中已安装 Python>=3.8.0PyTorch>=1.8。预训练模型和所需数据集将自动从最新的 YOLOv5 版本下载。

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

层冻结的工作原理#

神经网络中冻结层时,可以阻止其参数(权重和偏置)在训练过程中更新。在 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 架构通常由一个骨干网络(标准配置中的第 0-9 层,例如 YOLOv5s/m/l/x)和一个检测头(其余层)组成。骨干网络负责特征提取,检测头则执行目标检测

# 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 参数控制冻结哪些层。该参数指定第一个未冻结模块的索引;此索引之前的所有模块都会冻结其权重。如果需要确认哪些索引对应特定模块,可以使用 model.model(一个 nn.Sequential)检查模块顺序。

仅冻结骨干网络#

要冻结整个骨干网络(第 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

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

性能对比#

为了展示冻结层的影响,我们使用官方 COCO 预训练权重yolov5m.pt)对 YOLOv5m 在 Pascal VOC 数据集上训练了 50 个周期。我们比较了三种情形:训练所有层(--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 data/hyps/hyp.VOC.yaml --freeze 10

准确率结果#

结果表明,冻结层可以显著加快训练速度,但可能导致最终mAP(平均精度)略有下降。训练所有层通常能获得最佳准确率,而冻结更多层可以加快训练,但可能会牺牲性能。

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

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

YOLOv5 frozen layer training performance *Summary table of performance metrics*

资源利用率#

冻结更多层可以显著降低GPU内存需求和总体利用率。因此,在硬件资源有限的情况下,使用冻结层进行迁移学习非常有吸引力,你可以训练更大的模型或使用更大的图像尺寸,而这些通常可能难以实现。

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

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

何时使用层冻结#

在以下几种情况下,迁移学习期间冻结层尤其有利:

  1. 计算资源有限:当你的 GPU 内存或处理能力受限时。
  2. 数据集较小:当目标数据集明显小于原始预训练数据集时,冻结层有助于防止过拟合
  3. 快速原型开发:当你需要将现有模型快速适配到新的任务或领域以进行初步评估时。
  4. 特征领域相似:当新数据集中的低级特征与模型预训练所用数据集中的特征非常相似时。

在我们的术语表条目中进一步了解迁移学习的细节,并考虑使用超参数调优等技术来优化性能。

支持的环境#

Ultralytics 提供各种开箱即用的环境,其中预先安装了 CUDACuDNNPythonPyTorch 等必要依赖项。

项目状态#

YOLOv5 Continuous Integration Status

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

评论