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

引言

HRNet（High-Resolution Network）是一种用于语义分割任务的深度学习架构，其核心思想是在整个网络中保持高分辨率表示，而不是像传统方法那样先降低分辨率再恢复。本文将详细解析HRNet在COCO-Stuff数据集上的训练配置文件，帮助读者理解如何配置和优化HRNet模型进行语义分割任务。

硬件与基础配置

配置文件中首先定义了硬件相关的参数：

CUDNN:
  BENCHMARK: true
  DETERMINISTIC: false
  ENABLED: true
GPUS: (0,1,2,3)
WORKERS: 4

• CUDNN配置：启用了cuDNN加速，并设置了基准测试模式（BENCHMARK），这可以让cuDNN自动寻找最优的卷积算法实现
• GPU设置：使用了4块GPU进行并行训练（编号0-3）
• WORKERS：数据加载时使用4个子进程，可以加快数据预处理速度

数据集配置

DATASET:
  DATASET: cocostuff
  ROOT: '../../../../dataset/coco_stuff_10k/'
  TEST_SET: 'val.lst'
  TRAIN_SET: 'train.lst'
  NUM_CLASSES: 171

• COCO-Stuff数据集：这是一个包含171个类别的语义分割数据集，扩展自COCO数据集，增加了更多"stuff"类别的标注
• 数据路径：训练集和验证集分别通过train.lst和val.lst文件指定
• 类别数：171个语义类别需要模型进行区分

模型架构

HRNet的核心在于其多分辨率并行处理的结构：

MODEL:
  NAME: seg_hrnet
  PRETRAINED: 'pretrained_models/hrnetv2_w48_imagenet_pretrained.pth'
  EXTRA:
    STAGE1:
      NUM_MODULES: 1
      NUM_BRANCHES: 1
      BLOCK: BOTTLENECK
      NUM_BLOCKS: [4]
      NUM_CHANNELS: [64]
    STAGE2:
      NUM_BRANCHES: 2
      NUM_CHANNELS: [48, 96]
    STAGE3:
      NUM_MODULES: 4
      NUM_BRANCHES: 3
      NUM_CHANNELS: [48, 96, 192]
    STAGE4:
      NUM_MODULES: 3
      NUM_BRANCHES: 4
      NUM_CHANNELS: [48, 96, 192, 384]

• 网络结构：HRNet包含4个阶段(stage)，随着阶段推进，分支(分辨率)数量逐渐增加
• 多分辨率处理：从stage2开始，网络并行处理多个分辨率特征，并通过SUM方法进行特征融合
• 预训练权重：使用了在ImageNet上预训练的HRNet-W48模型作为基础

损失函数与OHEM

LOSS:
  USE_OHEM: true
  OHEMTHRES: 0.9
  OHEMKEEP: 131072

• OHEM(Online Hard Example Mining)：一种困难样本挖掘技术
  • 只保留损失值在前90%(OHEMTHRES)的样本进行反向传播
  • 每次保留131072个像素点(OHEMKEEP)用于计算梯度
  • 这种技术可以关注难以分类的样本，提高模型在困难区域的性能

训练策略

TRAIN:
  BATCH_SIZE_PER_GPU: 4
  TOTAL_BATCH_SIZE: 16  # 4 GPUs x 4
  BASE_SIZE: 520
  IMAGE_SIZE: [520, 520]
  BEGIN_EPOCH: 0
  END_EPOCH: 110
  OPTIMIZER: sgd
  LR: 0.001
  WD: 0.0001
  MOMENTUM: 0.9
  FLIP: true
  MULTI_SCALE: true

• 批量大小：每块GPU处理4张图像，4块GPU总批量大小为16
• 输入尺寸：图像被调整到520×520分辨率
• 训练周期：共训练110个epoch
• 优化器：使用带动量的SGD优化器
  • 学习率(LR)：0.001
  • 权重衰减(WD)：0.0001
  • 动量(MOMENTUM)：0.9
• 数据增强：
  • 随机水平翻转(FLIP)
  • 多尺度训练(MULTI_SCALE)，增强模型对不同尺寸物体的识别能力

测试配置

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

• 测试批量：每GPU处理1张图像
• 测试尺寸：固定520×520分辨率
• 测试时增强：未使用翻转和多尺度测试

关键参数解析

1. NONBACKBONE参数：

   NONBACKBONE_KEYWORDS: ['last_layer']
   NONBACKBONE_MULT: 10
   

   • 对非骨干网络部分（如最后的分类层）使用10倍的学习率
   • 这种策略常见于迁移学习，可以更快调整新添加的层

2. SCALE_FACTOR：

   SCALE_FACTOR: 16
   

   • 这是HRNet的输出步长，表示特征图相对于输入图像的缩小倍数
   • HRNet通过保持高分辨率，相比传统方法(如输出步长32)能保留更多细节

3. IGNORE_LABEL：

   IGNORE_LABEL: 255
   

   • 在损失计算时忽略标签值为255的像素
   • 这些通常是图像边界或无效区域

总结

这份配置文件展示了一个完整的HRNet-W48模型在COCO-Stuff数据集上的训练方案。关键点包括：

1. 使用多分辨率并行的HRNet架构，保持高分辨率特征
2. 采用OHEM策略关注困难样本
3. 使用适度的数据增强和多GPU训练
4. 对骨干网络和分类头采用不同的学习率

这种配置平衡了训练效率和模型性能，适合处理COCO-Stuff这类包含大量类别的复杂语义分割任务。通过调整这些参数，可以进一步优化模型在不同场景下的表现。