在 GitHub 上编辑页面 (EN)

YOLOv5 中冻结层的迁移学习

📚 This guide explains how to freeze YOLOv5 🚀 layers when implementing transfer learning. Transfer learning is a powerful machine learning (ML) technique that allows you to quickly retrain a model on new data without retraining the entire network from scratch. By freezing the weights of initial layers and only updating the parameters of later layers, you can significantly reduce computational resource requirements and training time. However, this approach might slightly impact the final model accuracy.

在你开始之前

首先，克隆 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 属性设置为 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 层）负责 特征提取，以及一个执行 目标检测 的头部（剩余层）。

# 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

当主要需要为不同数量的输出类别调整模型，同时保持绝大多数已学习的特征不变时，这种方法很有用。它对于 微调 (fine-tuning) 所需的计算资源最少。

性能对比

为了说明冻结层的效果，我们在 Pascal VOC 数据集 上对 YOLOv5m 进行了 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 (mean Average Precision) 略有下降。训练所有层通常能获得最佳准确率，而冻结更多层则以可能较低的性能为代价提供了更快的训练速度。

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

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

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

资源利用率

冻结更多层可以大幅降低 GPU 内存需求和整体利用率。这使得在处理有限硬件资源时，使用冻结层的迁移学习成为一个有吸引力的选择，允许训练更大的模型或使用比原本可能实现的更大的图像尺寸。

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

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

何时使用层冻结

在迁移学习过程中进行层冻结在以下几种情况下特别有利：

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

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

支持的环境

Ultralytics 提供了各种随时可用的环境，预装了 CUDA、CuDNN、Python 和 PyTorch 等基本依赖项。

项目状态

YOLOv5 持续集成状态

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

Contributors
在 GitHub 上编辑页面 (EN)
GLglenn-jocher16PDpderrenger1RIRizwanMunawar1

评论