YOLOv5 中冻结层的迁移学习#
📚 本指南解释了在实现迁移学习时如何冻结 YOLOv5 🚀 层。迁移学习是一种强大的机器学习 (ML) 技术,允许你快速在新的数据上重新训练模型,而无需从头开始重新训练整个网络。通过冻结初始层的权重并仅更新后期层的参数,你可以显著减少计算资源需求和训练时间。然而,这种方法可能会对最终模型的准确率产生轻微影响。
开始之前#
首先,克隆 YOLOv5 仓库并安装 requirements.txt 中列出的必要依赖项。确保你拥有安装了 PyTorch>=1.8 的 Python>=3.8.0 环境。预训练模型和所需的数据集将从最新的 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 range(freeze)] # Define layers to freeze based on module index
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 代码片段查看所有模块及其参数的名称:
# Assuming 'model' is your loaded YOLOv5 model instance
for name, param in model.named_parameters():
print(name)
"""
Example 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 架构通常由负责特征提取的主干网络(标准配置如 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 参数来控制哪些层被冻结。此参数指定第一个未冻结模块的索引;在此索引之前的所有模块的权重都将被冻结。如果需要确认哪些索引对应于特定区块,请使用 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当你主要需要针对不同的输出类别数量调整模型,同时保持绝大部分已学习特征不变时,这种方法非常有用。它对微调所需的计算资源最少。
性能比较#
为了说明冻结层的影响,我们在Pascal VOC dataset上对 YOLOv5m 进行了 50 个轮次的训练,从官方的 COCO 预训练权重(yolov5m.pt)开始。我们比较了三种场景:训练所有层(--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 对比
*Summary table of performance metrics*
资源利用率#
冻结更多层可大幅减少 GPU 内存需求和整体利用率。这使得在使用有限的硬件资源时,带有冻结层的迁移学习成为一个极具吸引力的选项,从而能够训练更大的模型或使用比平时更大的图像尺寸。
GPU 已分配内存百分比 (%)
GPU 利用率 (%)
何时使用层冻结#
在以下几种情况下,迁移学习期间的层冻结特别有优势:
- 计算资源有限:如果你对 GPU 内存或处理能力有严格限制。
- 小型数据集:当你的目标数据集明显小于原始预训练数据集时,冻结有助于防止过拟合。
- 快速原型设计:当你需要快速将现有模型调整到新任务或领域以进行初步评估时。
- 相似特征域:如果你的新数据集中的低级特征与模型预训练时所用的数据集非常相似。
在我们的词汇表条目中探索更多关于迁移 learning 微妙之处的内容,并考虑诸如超参数调优之类的方法来优化性能。
支持的环境#
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 小时以及每次新的代码提交时自动运行。