使用 YOLOv5 冻结层进行迁移学习#
📚 本指南介绍在实现迁移学习时,如何冻结YOLOv5 🚀 层。迁移学习是一种强大的机器学习(ML)技术,可以让你在新数据上快速重新训练模型,而无需从头开始重新训练整个网络。冻结初始层的权重,只更新后续层的参数,可以显著降低计算资源需求和训练时间。不过,这种方法可能会略微影响最终模型的准确率。
开始之前#
首先,克隆 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 (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
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 代码片段查看所有模块及其参数的名称:
# 假设“model”是已加载的 YOLOv5 模型实例
for name, param in model.named_parameters():
print(name)
"""
Example Output:
model.0.conv.weight
model.0.bn.weight
model.0.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 架构通常由骨干网络(标准配置如 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 参数来控制冻结哪些层。此参数指定第一个未冻结模块的索引;该索引之前的所有模块都会被冻结权重。也可以传入多个索引(例如 --freeze 0 1 2),这样只会冻结指定的模块。如果需要确认特定模块对应哪些索引,可以使用 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当你主要需要调整模型以适应不同数量的输出类别,同时保留绝大部分已学习特征时,这种方法很有用。它只需为微调消耗极少的计算资源。
性能对比#
为了说明冻结层的效果,我们使用官方 COCO 预训练权重(yolov5m.pt),在 Pascal VOC 数据集上训练了 YOLOv5m 50 个轮次。我们比较了三种情形:训练所有层(--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 data/hyps/hyp.VOC.yaml --freeze 10准确率结果#
结果表明,冻结层可以显著加快训练速度,但可能会略微降低最终的 mAP(平均精度均值)。训练所有层通常能取得最佳准确率,而冻结更多层可以加快训练,但可能会降低性能。
mAP50 训练对比
mAP50-95 训练对比
*性能指标汇总表*
资源利用情况#
冻结更多层可以大幅降低 GPU 内存需求和整体资源利用率。因此,在硬件资源有限时,使用冻结层进行迁移学习很有吸引力;这样你就能训练更大的模型,或使用更大的图像尺寸。
GPU 内存分配(%)
GPU 利用率(%)
何时使用层冻结#
在以下几种情况下,迁移学习期间冻结层尤其有优势:
- 计算资源有限:如果你的 GPU 内存或处理能力有限。
- 数据集较小:当目标数据集远小于原始预训练数据集时,冻结层有助于防止过拟合。
- 快速原型验证:当你需要快速将现有模型适配到新任务或新领域,以便进行初步评估时。
- 特征领域相似:如果新数据集中的低级特征与模型预训练所用数据集中的特征非常相似。
在我们的术语表条目中了解更多迁移学习的细节,并考虑使用超参数进化来优化性能。
支持的环境#
Ultralytics 提供多种开箱即用的环境,其中已预装 CUDA、CuDNN、Python 和 PyTorch 等必要依赖项。
- 免费 GPU 笔记本:
- Google Cloud:GCP 快速入门指南
- Amazon:AWS 快速入门指南
- Azure:AzureML 快速入门指南
- Docker:Docker 快速入门指南
项目状态#
此徽章表示所有 YOLOv5 GitHub Actions 持续集成(CI)测试均已通过。这些 CI 测试会严格评估 YOLOv5 在关键操作中的功能和性能,包括训练、验证、推理、导出和基准测试。测试确保 YOLOv5 在 macOS、Windows 和 Ubuntu 上稳定可靠地运行,并且每 24 小时自动运行一次,也会在每次提交新代码时运行。
常见问题#
它会冻结
model.model的前N个模块,使这些模块的权重保持不变且不接收梯度。--freeze 10会冻结骨干网络,而--freeze 24会冻结除最终Detect模块以外的所有部分。在上面的 VOC 实验中,训练所有层获得了最佳 mAP;冻结骨干网络会略微降低准确率;除检测头外全部冻结则会造成最大的准确率损失。作为回报,冻结层的训练速度更快,且 GPU 内存占用大幅降低。
当你的图像与预训练数据共享低级特征、数据集较小,或者你需要在硬件有限的情况下快速获得基线结果时,都适合冻结骨干网络。数据集扩大后,或当你需要进一步提升 mAP 时,可以训练所有层。
可以。使用冻结模型的
best.pt启动新的训练运行,并移除--freeze,即可对整个网络进行微调,通常应使用较低的学习率。