冻结层迁移学习

📚 本指南解释了如何在迁移学习时冻结YOLOv5 🚀 层。迁移学习是一种有效的方法，可以在不重新训练整个网络的情况下，快速在新数据上重新训练模型。相反，初始权重的一部分被冻结，其余权重用于计算损失并通过优化器更新。这比正常训练需要更少的资源，并允许更快的训练时间，尽管它也可能导致最终训练精度的降低。

开始之前

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

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

冻结骨干网络

在训练开始前，train.py中freeze列表中匹配的所有层将被冻结，即将它们的梯度设置为零。

# 冻结
freeze = [f"model.{x}." for x in range(freeze)]  # 要冻结的层
for k, v in model.named_parameters():
    v.requires_grad = True  # 训练所有层
    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)

"""输出:
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:
    # [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:
    - [-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.' - 'model.9.'的所有模块：

python train.py --freeze 10

冻结所有层

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

python train.py --freeze 24

结果

我们在上述两种场景下训练YOLOv5m在VOC数据集上，以及一个默认模型（无冻结），从官方COCO预训练的--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利用率也越低。这表明较大的模型，或者在更大的--image-size下训练的模型，可能会从冻结中受益，以加快训练速度。

训练GPU内存分配百分比

训练GPU内存利用率百分比

支持的环境

Ultralytics 提供了一系列开箱即用的环境，每个环境都预装了必要的依赖项，如 CUDA、CUDNN、Python 和 PyTorch，以启动您的项目。

免费GPU笔记本:
Google Cloud: GCP快速入门指南
Amazon: AWS快速入门指南
Azure: AzureML快速入门指南
Docker: Docker快速入门指南

项目状态

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