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

概述

HRNet（High-Resolution Network）是一种用于语义分割任务的深度学习网络架构，其核心特点是能够在整个网络中保持高分辨率表示，而不是像传统方法那样先降低分辨率再恢复。本文将以HRNet-Semantic-Segmentation项目中针对COCO-Stuff数据集的配置文件为例，详细解析HRNet-W48模型的训练参数设置和技术细节。

硬件与基础配置

配置文件中首先定义了基础的硬件和运行环境设置：

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

• CUDNN配置：启用了cuDNN加速（ENABLED: true），并设置了BENCHMARK模式，这可以让cuDNN自动寻找最优的卷积算法实现。DETERMINISTIC设为false以获得更好的性能。
• GPU设置：使用4块GPU进行训练（索引0-3），充分利用多GPU并行计算能力。
• 工作线程：设置4个数据加载工作线程（WORKERS: 4），平衡数据加载速度和内存消耗。
• 日志输出：每10个批次打印一次训练信息（PRINT_FREQ: 10）。

数据集配置

DATASET:
  DATASET: cocostuff
  ROOT: 'data/'
  TEST_SET: 'list/cocostuff/val.lst'
  TRAIN_SET: 'list/cocostuff/train.lst'
  NUM_CLASSES: 171

• 数据集选择：使用COCO-Stuff数据集，这是一个包含171个语义类别的复杂场景理解数据集。
• 数据路径：训练集和验证集的列表文件分别指定，便于管理大规模数据集。
• 类别数量：明确指定171个类别，这对于网络最后一层的设计至关重要。

模型架构

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

HRNet-W48模型的核心特点：

1. 多分辨率并行处理：网络包含四个阶段，每个阶段都有多个分支处理不同分辨率的特征图。

   • 阶段1：单分支，64通道
   • 阶段2：双分支，48和96通道
   • 阶段3：三分支，48、96和192通道
   • 阶段4：四分支，48、96、192和384通道

2. 特征融合：使用SUM方法融合不同分支的特征，保持信息流动。

3. 预训练权重：加载在ImageNet上预训练的HRNetV2-W48权重，这是提升模型性能的关键。

4. 最终卷积：使用1x1卷积（FINAL_CONV_KERNEL: 1）将特征映射到类别空间。

损失函数与训练策略

LOSS:
  USE_OHEM: false
  OHEMTHRES: 0.9
  OHEMKEEP: 131072

TRAIN:
  IMAGE_SIZE: [520,520]
  BASE_SIZE: 520
  BATCH_SIZE_PER_GPU: 4
  SHUFFLE: true
  BEGIN_EPOCH: 0
  END_EPOCH: 110
  RESUME: true
  OPTIMIZER: sgd
  LR: 0.001
  WD: 0.0001
  NONBACKBONE_KEYWORDS: ['last_layer']
  NONBACKBONE_MULT: 10
  MOMENTUM: 0.9
  NESTEROV: false
  FLIP: true
  MULTI_SCALE: true
  DOWNSAMPLERATE: 1
  IGNORE_LABEL: 255
  SCALE_FACTOR: 16

训练细节

1. 图像尺寸：使用520x520的输入分辨率，这是COCO-Stuff数据集的常见处理尺寸。

2. 批量大小：每GPU 4张图像，4块GPU总批量大小为16，平衡了内存使用和训练稳定性。

3. 优化器：使用SGD优化器，学习率0.001，权重衰减0.0001，动量0.9。

   • 对非骨干网络部分（特别是最后一层）使用10倍学习率（NONBACKBONE_MULT: 10），加速特定任务的适应。

4. 数据增强：

   • 随机水平翻转（FLIP: true）
   • 多尺度训练（MULTI_SCALE: true）
   • 忽略标签值为255的像素（通常用于标记难样本或边界区域）

5. 训练周期：共110个epoch，考虑到COCO-Stuff数据集的规模和复杂性，这个周期数足够让模型收敛。

6. 损失函数：未使用OHEM（Online Hard Example Mining），而是采用标准的交叉熵损失。

测试配置

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

测试阶段配置特点：

1. 批量大小：每GPU 1张图像，确保评估过程的稳定性。

2. 测试增强：默认关闭多尺度和翻转测试，保持评估的一致性。

3. 样本数量：指定测试200个样本，便于快速验证模型性能。

技术要点总结

1. HRNet架构优势：通过并行多分辨率分支和重复的多尺度融合，HRNet能够在整个网络中保持高分辨率表示，这对语义分割任务尤为重要。

2. 训练技巧：

   • 使用ImageNet预训练权重加速收敛
   • 对不同网络部分采用差异化的学习率
   • 综合使用多种数据增强手段

3. 参数选择：

   • 学习率和权重衰减经过精心调整
   • 批量大小与GPU数量匹配
   • 训练周期足够长以确保充分收敛

这份配置文件展示了如何在复杂语义分割任务上配置HRNet模型，其中的参数选择和架构设计都经过了充分验证，可以作为类似任务的参考模板。理解这些配置背后的原理，有助于研究人员根据自身需求进行调整和优化。