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

概述

本文将深入解析HRNet语义分割项目中针对ADE20K数据集的训练配置文件，帮助读者理解如何配置一个高性能的语义分割模型训练过程。HRNet（High-Resolution Network）是一种保持高分辨率特征表示的网络结构，在语义分割任务中表现出色。

硬件与基础配置

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

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

• CUDNN配置：启用了cuDNN加速，并开启基准测试模式以选择最优算法
• GPU设置：使用了8块GPU进行分布式训练
• WORKERS：数据加载使用8个工作进程，可提高数据吞吐量

数据集配置

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

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

关键点：

• 训练集和验证集分别通过train.lst和val.lst指定
• 类别数为150，与ADE20K数据集一致
• 数据集路径需要根据实际存放位置调整

模型架构

HRNet-OCR模型配置如下：

MODEL:
  NAME: seg_hrnet_ocr
  NUM_OUTPUTS: 2
  PRETRAINED: '../../../../dataset/pretrained_models/HRNet_W48_C_ssld_pretrained.pth'

• seg_hrnet_ocr：结合了HRNet和OCR(Object Contextual Representation)的分割模型
• 预训练权重：使用了在ImageNet上预训练的HRNet-W48模型
• NUM_OUTPUTS：设置为2，可能用于多任务学习或辅助损失

网络结构细节

HRNet采用多分支结构，逐步增加分支数量：

EXTRA:
  STAGE1:
    NUM_MODULES: 1
    NUM_RANCHES: 1
    BLOCK: BOTTLENECK
    NUM_BLOCKS: [4]
    NUM_CHANNELS: [64]
  STAGE2:
    NUM_BRANCHES: 2
    NUM_CHANNELS: [48, 96]
  STAGE3:
    NUM_BRANCHES: 3
    NUM_CHANNELS: [48, 96, 192]
  STAGE4:
    NUM_BRANCHES: 4
    NUM_CHANNELS: [48, 96, 192, 384]

网络特点：

1. 四阶段结构，分支数从1增加到4
2. 使用BOTTLENECK块降低第一阶段计算量
3. 后续阶段使用BASIC块
4. 通道数逐步增加，保持高分辨率表示

损失函数配置

LOSS:
  USE_OHEM: true
  OHEMTHRES: 0.9
  OHEMKEEP: 131072
  BALANCE_WEIGHTS: [0.4, 1]

• OHEM（Online Hard Example Mining）：在线难例挖掘，专注于难以分类的样本
• OHEMTHRES：0.9表示只选择损失值在前90%的样本
• OHEMKEEP：保留131072个像素点用于计算损失
• 权重平衡：辅助损失和主损失的权重比为0.4:1

训练策略

TRAIN:
  IMAGE_SIZE: [520, 520]
  BATCH_SIZE_PER_GPU: 2
  BATCH_SIZE: 16  # 8GPUs×2
  BEGIN_EPOCH: 0
  END_EPOCH: 120
  OPTIMIZER: sgd
  LR: 0.02
  WD: 0.0001
  MOMENTUM: 0.9

训练关键参数：

• 输入尺寸：520×520像素
• 批量大小：每个GPU处理2张图像，8GPU共16张
• 训练周期：120个epoch
• 优化器：使用带动量的SGD
• 学习率：初始为0.02，配合多项式衰减
• 权重衰减：L2正则化系数为0.0001
• 数据增强：启用多尺度训练和随机翻转

验证配置

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

验证时：

• 使用单尺度测试（520×520）
• 不启用测试时数据增强（如多尺度、翻转）
• 每GPU处理1张图像

实际应用建议

1. 硬件适配：根据实际GPU数量调整BATCH_SIZE_PER_GPU，保持总batch size为16左右
2. 学习率调整：如果改变batch size，需按线性规则调整学习率
3. 内存优化：520×520分辨率需要较大显存，可适当降低分辨率或batch size
4. 训练监控：利用PRINT_FREQ=10每10个iteration打印一次日志

通过这份配置文件，我们可以看到一个完整的HRNet-OCR语义分割模型的训练设置，从数据准备到模型架构，再到训练策略，每个环节都经过精心设计。理解这些配置参数对于调整模型性能、适应不同任务需求至关重要。