跳至内容

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

本指南介绍了在实施迁移学习时如何冻结 YOLOv5 🚀 层。迁移学习是一种功能强大的技术,它可以让您在新数据上快速重新训练模型,而无需重新训练整个网络。通过冻结部分初始权重,只更新其余权重,可以大大减少计算资源和训练时间,不过这种方法可能会稍微影响最终模型的准确性。

开始之前

克隆该 repo 并将requirements.txt安装在 环境中安装 requirements.txt环境中安装 requirements.txt,包括 PyTorch.8.模型数据集会自动从最新的YOLOv5 版本下载。

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

冻结层的工作原理

当你冻结神经网络中的层时,基本上就是将它们的参数设置为不可训练。这些层的梯度被设置为零,从而防止在反向传播过程中出现任何权重更新。这在YOLOv5 的训练过程中实现如下:

# Freeze
freeze = [f"model.{x}." for x in range(freeze)]  # layers to freeze
for k, v in model.named_parameters():
    v.requires_grad = True  # train all layers
    if any(x in k for x in freeze):
        print(f"freezing {k}")
        v.requires_grad = False

探索模型架构

要有效冻结模型的特定部分,了解层结构很有帮助。您可以使用

for k, v in model.named_parameters():
    print(k)

"""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 层)和头部(其余各层)组成:

# YOLOv5 v6.0 backbone
backbone:
    # [from, number, module, args]
    - [-1, 1, Conv, [64, 6, 2, 2]] # 0-P1/2
    - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
    - [-1, 3, C3, [128]]
    - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
    - [-1, 6, C3, [256]]
    - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
    - [-1, 9, C3, [512]]
    - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
    - [-1, 3, C3, [1024]]
    - [-1, 1, SPPF, [1024, 5]] # 9

# YOLOv5 v6.0 head
head:
    - [-1, 1, Conv, [512, 1, 1]]
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
    - [[-1, 6], 1, Concat, [1]] # cat backbone P4
    - [-1, 3, C3, [512, False]] # 13
    # ... remaining head layers

冷冻选项

仅冻结主干网

只冻结骨干层(0-9 层),这有利于模型适应新的类别,同时保留已学的特征提取功能:

python train.py --freeze 10

当新数据集与原始训练数据具有相似的底层特征,但类别或对象不同时,这种方法尤为有效。

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

冻结整个模型,检测模块中的最终输出卷积层除外:

python train.py --freeze 24

当你想保留模型的大部分已学特征,但又需要对其进行调整以检测不同数量的类别时,这种方法是理想的选择。

性能比较

我们从官方 COCO 预训练权重开始,使用不同的冻结策略在 VOC 数据集上训练 YOLOv5m:

python train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp hyp.finetune.yaml

准确性结果

结果表明,冻结层可以加快训练速度,但会略微降低最终准确率

冷冻训练 mAP50 结果

冷冻训练 mAP50-95 结果

表格结果

资源利用

冻结更多图层可降低GPU 内存需求和利用率,因此这项技术对于训练更大的模型或使用更高分辨率的图像非常有价值:

培训GPU 内存分配百分比

培训GPU 内存利用率百分比

何时使用层冷冻

迁移学习中的层冻结在以下情况下尤为有效:

  1. 计算资源有限:当GPU 内存或处理能力有限时
  2. 小数据集:当新数据集太小,无法在不过度拟合的情况下训练完整模型时
  3. 快速适应:当您需要快速调整模型以适应新领域时
  4. 针对特定任务进行微调:将通用模型调整为专门应用时

有关迁移学习技术及其应用的更多信息,请参阅迁移学习词汇条目

支持的环境

Ultralytics 提供了一系列随时可用的环境,每个环境都预装了基本的依赖项,如 CUDACUDNNPythonPyTorch等基本依赖项,以便启动项目。

项目现状

YOLOv5 CI

此徽章表示YOLOv5 GitHub Actions 的所有持续集成(CI)测试均已成功通过。这些 CI 测试严格检查了YOLOv5 在训练验证推理导出基准等多个关键方面的功能和性能。它们确保在 macOS、Windows 和 Ubuntu 上运行的一致性和可靠性,每 24 小时和每次新提交时都会进行一次测试。

📅创建于 1 年前 ✏️已更新 7 天前

评论