跳至内容

冻结层的迁移学习

📚 This guide explains how to freeze YOLOv5 🚀 layers when transfer learning. Transfer learning is a useful way to quickly retrain a model on new data without having to retrain the entire network. Instead, part of the initial weights are frozen in place, and the rest of the weights are used to compute loss and are updated by the optimizer. This requires less resources than normal training and allows for faster training times, though it may also result in reductions to final trained accuracy.

开始之前

克隆 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

精度比较

The results show that freezing speeds up training, but reduces final accuracy slightly.

冷冻训练 mAP50 结果

冷冻训练 mAP50-95 结果

表格结果

GPU 使用情况比较

有趣的是,冻结的模块越多,训练所需的GPU 内存就越少,GPU 利用率就越低。这表明,较大的模型或以较大图像大小训练的模型可能会受益于冻结,以便更快地训练。

培训GPU 内存分配百分比

培训GPU 内存利用率百分比

支持的环境

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

项目现状

YOLOv5 CI

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

📅 Created 11 months ago ✏️ Updated 5 days ago

评论