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

概述

本文将深入解析HRNet语义分割项目在ADE20K数据集上的训练配置文件，帮助读者理解如何配置和优化HRNet模型进行语义分割任务。ADE20K是一个包含150个类别的场景解析数据集，广泛应用于语义分割研究。

硬件与基础配置

CUDNN:
  BENCHMARK: true
  DETERMINISTIC: false
  ENABLED: true
GPUS: (0,1,2,3,4,5,6,7)
OUTPUT_DIR: 'output'
LOG_DIR: 'log'
WORKERS: 8
PRINT_FREQ: 10

这部分配置定义了训练的基础环境：

• CUDNN设置：启用cuDNN加速，但不使用确定性算法以保证性能
• GPU配置：使用8块GPU进行并行训练
• 工作线程：8个数据加载线程确保数据供给充足
• 日志输出：每10个batch打印一次训练信息

数据集配置

DATASET:
  DATASET: ade20k
  ROOT: '../../../../dataset/ade20k/'
  TEST_SET: 'val.lst'
  TRAIN_SET: 'train.lst'
  NUM_CLASSES: 150

关键点解析：

• ADE20K数据集包含150个语义类别
• 训练集和验证集通过列表文件(train.lst/val.lst)指定
• 路径配置需要注意相对路径的正确性

模型架构

MODEL:
  NAME: seg_hrnet
  NUM_OUTPUTS: 1
  PRETRAINED: '../../../../dataset/pretrained_models/HRNet_W48_C_ssld_pretrained.pth'
  EXTRA:
    FINAL_CONV_KERNEL: 1
    STAGE1-4: ...

HRNet模型特点：

1. 多分辨率特征融合：通过4个stage逐步增加分支数量(1→2→3→4)
2. 通道数配置：48/96/192/384的通道设计平衡了计算量和特征表达能力
3. 预训练权重：使用HRNet-W48在ImageNet上的半监督学习(SSLD)预训练模型
4. 最终卷积：使用1×1卷积将特征映射到分类空间

损失函数配置

LOSS:
  USE_OHEM: true
  OHEMTHRES: 0.9
  OHEMKEEP: 131072

OHEM(Online Hard Example Mining)策略：

• 作用：专注于难以分类的样本，提升模型在困难区域的表现
• 阈值：0.9表示只考虑预测置信度低于90%的样本
• 保留数量：131072个困难样本参与每批次训练

训练策略

TRAIN:
  IMAGE_SIZE: [520, 520]
  BASE_SIZE: 520
  BATCH_SIZE_PER_GPU: 2
  BEGIN_EPOCH: 0
  END_EPOCH: 120
  OPTIMIZER: sgd
  LR: 0.02
  WD: 0.0001
  MOMENTUM: 0.9
  FLIP: true
  MULTI_SCALE: true

训练关键参数：

1. 优化器：SGD+momentum(0.9)，学习率0.02，权重衰减0.0001
2. 数据增强：
   • 随机水平翻转(FLIP)
   • 多尺度训练(MULTI_SCALE)
3. 训练周期：120个epoch
4. 批大小：每GPU 2张图像，8GPU合计16张

测试配置

TEST:
  IMAGE_SIZE: [520, 520]
  BATCH_SIZE_PER_GPU: 1
  FLIP_TEST: false
  MULTI_SCALE: false

测试阶段注意事项：

• 关闭数据增强(FLIP_TEST/MULTI_SCALE=false)
• 使用固定520×520分辨率评估
• 每GPU批大小设为1，确保评估精度

实际应用建议

1. 硬件适配：GPU数量不足时可调整BATCH_SIZE_PER_GPU，但需相应调整学习率
2. 学习率策略：可考虑添加学习率衰减进一步提升模型性能
3. 显存优化：520分辨率较大，显存不足时可适当降低分辨率或批大小
4. 训练监控：利用LOG_DIR记录的日志分析训练过程

通过这份配置文件，我们可以清晰地了解HRNet在ADE20K数据集上的完整训练流程和参数设置，为实际应用和调优提供了可靠的基础。