跳转至内容

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

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

开始之前

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

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

层冻结的工作原理

当你在一个模型中冻结层时 神经网络，您可以阻止它们的参数（权重和偏差）在训练过程中被更新。在 PyTorch 中，这可以通过设置 requires_grad 该层 tensor 的属性 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 架构通常由一个 backbone（主干网络）（在 YOLOv5s/m/l/x 等标准配置中为第 0-9 层）负责 特征提取，以及一个执行 目标检测 的 head（剩余层）组成。

# 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

当您主要需要调整模型以适应不同数量的输出类别，同时保持绝大多数已学习的特征不变时，此方法非常有用。它需要最少的计算资源来进行 微调。

性能对比

为了说明冻结层的影响，我们在 Pascal VOC 数据集 持续 50 个 epoch，从官方 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 利用率 (%)

何时使用层冻结

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

  1. 有限的计算资源: 如果您的 GPU 内存或处理能力受到限制。
  2. 小数据集: 当目标数据集明显小于原始预训练数据集时，冻结有助于防止过拟合。
  3. 快速原型设计：当您需要快速地将现有模型调整到新的任务或领域以进行初步评估时。
  4. 相似特征域: 如果新数据集中的底层特征与模型预训练所用的数据集中的特征非常相似。

在我们的术语表条目中了解更多关于迁移学习的细微之处，并考虑使用超参数调整等技术来优化性能。

支持的环境

Ultralytics 提供了多种开箱即用的环境，这些环境预装了必要的依赖项，如 CUDA、CuDNN、Python 和 PyTorch。

项目状态

YOLOv5 持续集成状态

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



📅 创建于 1 年前 ✏️ 更新于 4 个月前
glenn-jocher UltralyticsAssistant RizwanMunawar

评论