跳至内容

冻结层的迁移学习

本指南介绍了如何在迁移学习时冻结 YOLOv5 🚀 层。迁移学习是根据新数据快速重新训练模型的有效方法，而无需重新训练整个网络。在转移学习中，部分初始权重被冻结在原位，其余权重用于计算损失，并由优化器进行更新。这比普通训练所需的资源更少，训练时间也更短，但也可能导致最终训练精度的降低。

开始之前

克隆 repo 并将requirements.txt安装在 Python>=3.8.0环境中安装 requirements txt，包括 PyTorch>=1.8.模型和数据集会自动从最新的YOLOv5 版本下载。

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

冻结骨干网

与 train.py 匹配的所有图层 freeze 将在训练开始前通过将梯度设为零来冻结 train.py 中的列表。

# 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
...
model.23.m.0.cv2.bn.weight
model.23.m.0.cv2.bn.bias
model.24.m.0.weight
model.24.m.0.bias
model.24.m.1.weight
model.24.m.1.bias
model.24.m.2.weight
model.24.m.2.bias
"""

纵观模型结构，我们可以看到模型的主干是 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

    - [-1, 1, Conv, [256, 1, 1]]
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
    - [[-1, 4], 1, Concat, [1]] # cat backbone P3
    - [-1, 3, C3, [256, False]] # 17 (P3/8-small)

    - [-1, 1, Conv, [256, 3, 2]]
    - [[-1, 14], 1, Concat, [1]] # cat head P4
    - [-1, 3, C3, [512, False]] # 20 (P4/16-medium)

    - [-1, 1, Conv, [512, 3, 2]]
    - [[-1, 10], 1, Concat, [1]] # cat head P5
    - [-1, 3, C3, [1024, False]] # 23 (P5/32-large)

    - [[17, 20, 23], 1, Detect, [nc]] # Detect(P3, P4, P5)

因此，我们可以将冻结列表定义为包含所有名称中包含 "model.0.- 模型.9. "的模块：

python train.py --freeze 10

冻结所有图层

为了在 Detect() 中冻结除最终输出卷积层之外的整个模型，我们将冻结列表设置为包含所有名称中带有 "model.0.- 模型.23. "的所有模块：

python train.py --freeze 24

成果

我们在 VOC 上对 YOLOv5m 进行了上述两种情况的训练，同时还对默认模型（无冻结）进行了训练。 --weights yolov5m.pt:

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 利用率就越低。这表明，较大的模型或以较大图像大小训练的模型可能会受益于冻结，以便更快地训练。

培训GPU 内存分配百分比

培训GPU 内存利用率百分比

支持的环境

Ultralytics 提供了一系列随时可用的环境，每个环境都预装了基本的依赖项，如 CUDA、CUDNN、 Python和 PyTorch等基本依赖项，以便启动您的项目。

项目现状

YOLOv5 CI

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



创建于 2023-11-12，更新于 2024-06-10
作者:glenn-jocher(7)

评论