HRNet语义分割模型在ADE20K数据集上的训练配置解析

配置概述

本文详细解析HRNet语义分割项目中用于ADE20K数据集的训练配置文件seg_hrnet_w48_520x520_sgd_lr2e-2_wd1e-4_bs_16_epoch120.yaml。该配置定义了使用HRNet-W48架构在ADE20K数据集上进行语义分割训练的各项参数。

硬件与基础设置

配置首先定义了硬件相关的基础设置：

• CUDNN：启用cuDNN加速库，开启基准测试模式以提高性能，但不启用确定性模式以保持灵活性
• GPU设置：使用4块GPU(0-3)进行训练
• 目录设置：输出目录为'output'，日志目录为'log'
• 数据加载：使用4个工作进程(workers)加载数据，每10个批次打印一次训练信息

数据集配置

ADE20K是一个包含150个类别的场景解析数据集：

• 数据集路径：根目录为'data/'，训练集和验证集分别位于'list/ade20k/'目录下
• 类别数量：150个语义类别
• 图像尺寸：训练和测试时使用520x520的输入分辨率

模型架构

配置使用HRNet-W48作为基础网络：

• 预训练权重：加载在ImageNet上预训练的HRNetv2-W48模型
• 网络结构：
  • 包含4个stage的多分辨率并行结构
  • 各stage的分支数分别为1、2、3、4
  • 通道数从48逐渐增加到384
  • 使用BASIC和BOTTLENECK两种基础块
  • 特征融合方法采用简单的SUM操作

训练策略

训练过程采用以下优化策略：

• 优化器：使用SGD优化器
  • 初始学习率0.02
  • 权重衰减0.0001
  • 动量0.9
  • 不使用Nesterov动量
• 数据增强：
  • 随机水平翻转
  • 多尺度训练
  • 缩放因子为11
• 批次设置：每GPU批次大小为4，4块GPU总批次大小为16
• 训练周期：共120个epoch
• 损失函数：不使用OHEM(在线难例挖掘)

测试配置

测试阶段采用以下设置：

• 输入尺寸：保持520x520分辨率
• 批次大小：每GPU处理1张图像
• 测试样本数：200张
• 测试增强：不使用多尺度和翻转测试

技术要点解析

1. HRNet多分辨率特性：配置文件清晰地展示了HRNet的多分辨率并行结构，从stage1的单一分辨率逐渐增加到stage4的四种分辨率并行处理。

2. 训练技巧：

   • 较大的初始学习率(0.02)配合SGD优化器是该配置的特点
   • 适中的权重衰减(1e-4)有助于防止过拟合
   • 多尺度训练增强模型对不同尺寸目标的识别能力

3. ADE20K适配：配置专门针对ADE20K数据集的150个类别进行了优化，输入分辨率520x520在精度和计算成本间取得了良好平衡。

实践建议

1. 对于计算资源有限的用户，可以适当减少批次大小或使用较少的GPU
2. 训练初期可以尝试降低学习率观察收敛情况
3. 在显存允许的情况下，可以尝试增大输入尺寸以提高分割精度
4. 对于小目标较多的场景，可以适当增加多尺度训练中的缩放范围

通过这份配置文件，我们可以深入了解HRNet在语义分割任务上的实现细节和优化策略，为自定义训练提供了良好的参考模板。