Lite-HRNet：一种轻量级高分辨率网络

简介

这是 Lite-HRNet：一种轻量级高分辨率网络 的官方 PyTorch 实现。在本工作中，我们提出了一种用于人体姿态估计的高效高分辨率网络——Lite-HRNet。我们首先将 ShuffleNet 中的高效 Shuffle 块简单地应用于 HRNet（高分辨率网络），从而在性能上超越了诸如 MobileNet、ShuffleNet 和 Small HRNet 等流行的轻量级网络。然而，我们发现 Shuffle 块中广泛使用的逐点（1x1）卷积成为了计算瓶颈。为此，我们引入了一种轻量级单元——条件通道加权，以替代 Shuffle 块中昂贵的逐点（1x1）卷积。通道加权的复杂度与通道数呈线性关系，远低于逐点卷积的二次时间复杂度。我们的解决方案从所有通道以及 HRNet 并行分支中已有的多尺度特征中学习权重，并利用这些权重作为跨通道和跨尺度信息交换的桥梁，从而弥补了逐点卷积的作用。实验结果表明，Lite-HRNet 在人体姿态估计任务上优于现有的轻量级网络。此外，Lite-HRNet 还可以以同样的轻量化方式轻松应用于语义分割任务。

结果与模型

COCO val2017 数据集上的结果（检测器在 COCO val2017 上的人体 AP 为 56.4）

架构	输入尺寸	参数量	FLOPs	AP	AP50	AP75	AR	AR50	模型下载
Naive Lite-HRNet-18	256x192	0.7M	194.8M	0.628	0.855	0.699	0.691	0.901	GoogleDrive 或 OneDrive
Wider Naive Lite-HRNet-18	256x192	1.3M	311.1M	0.660	0.871	0.737	0.721	0.913	GoogleDrive 或 OneDrive
Lite-HRNet-18	256x192	1.1M	205.2M	0.648	0.867	0.730	0.712	0.911	GoogleDrive 或 OneDrive
Lite-HRNet-18	384x288	1.1M	461.6M	0.676	0.878	0.750	0.737	0.921	GoogleDrive 或 OneDrive
Lite-HRNet-30	256x192	1.8M	319.2M	0.672	0.880	0.750	0.733	0.922	GoogleDrive 或 OneDrive
Lite-HRNet-30	384x288	1.8M	717.8M	0.704	0.887	0.777	0.762	0.928	GoogleDrive 或 OneDrive

MPII 验证集上的结果

架构	输入尺寸	参数量	FLOPs	均值	均值@0.1	模型下载
Naive Lite-HRNet-18	256x256	0.7M	259.6M	0.853	0.305	GoogleDrive 或 OneDrive
Wider Naive Lite-HRNet-18	256x256	1.3M	418.7M	0.868	0.311	GoogleDrive 或 OneDrive
Lite-HRNet-18	256x256	1.1M	273.4M	0.854	0.295	GoogleDrive 或 OneDrive
Lite-HRNet-30	256x256	1.8M	425.3M	0.870	0.313	GoogleDrive 或 OneDrive

运行环境

代码使用 Python 3.6 在 Ubuntu 16.04 上开发。需要 NVIDIA GPU。代码是在 8 张 NVIDIA V100 GPU 卡上开发并测试的。其他平台或 GPU 卡尚未经过全面测试。

快速入门

环境要求

• Linux（Windows 尚未得到官方支持）
• Python 3.6+
• PyTorch 1.3+
• CUDA 9.2+（若从源码编译 PyTorch，CUDA 9.0 亦可兼容）
• GCC 5+
• mmcv（请安装最新版本的 mmcv-full）
• Numpy
• cv2
• json_tricks
• xtcocotools

安装

a. 安装 mmcv，我们建议您按照以下方式安装预编译的 mmcv。

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html

请将 URL 中的 {cu_version} 和 {torch_version} 替换为您所需的版本。例如，要安装带有 CUDA 11 和 PyTorch 1.7.0 的最新 mmcv-full，可以使用以下命令：

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu110/torch1.7.0/index.html

如果在安装过程中需要编译，请确保 CUDA 版本和 PyTorch 版本与 mmcv-full 安装命令中的版本完全一致。例如，PyTorch 1.7.0 和 1.7.1 被视为不同的版本。 有关不同 MMCV 版本与不同 PyTorch 和 CUDA 版本的兼容性，请参阅 这里。

可选地，您也可以通过以下命令从源码编译 mmcv：

git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
MMCV_WITH_OPS=1 pip install -e .  # 此步骤后将安装包含 CUDA 操作的 mmcv-full 包
# 或者 pip install -e .  # 此步骤后将安装不包含 CUDA 操作的 mmcv 包
cd ..

或者直接运行：

pip install mmcv-full
# 或者：pip install mmcv

重要提示： 如果您已经安装了 mmcv，请务必先运行 pip uninstall mmcv。如果同时安装了 mmcv 和 mmcv-full，将会出现 ModuleNotFoundError。

b. 安装构建依赖

pip install -r requirements.txt

准备数据集

建议将数据集根目录创建符号链接到 $LITE_HRNET/data。 如果您的文件夹结构不同，可能需要修改配置文件中的相应路径。

对于 COCO 数据集，请从 COCO 下载下载，训练和验证 COCO 关键点任务时需要使用 2017 年的训练集和验证集。HRNet-Human-Pose-Estimation 提供了 COCO val2017 的人员检测结果，用于复现我们的多人姿态估计结果。请从 OneDrive 下载。 将下载的文件解压到 $LITE_HRNET/data 目录下，使其呈现如下结构：

lite_hrnet
├── configs
├── models
├── tools
`── data
    │── coco
        │-- annotations
        │   │-- person_keypoints_train2017.json
        │   |-- person_keypoints_val2017.json
        |-- person_detection_results
        |   |-- COCO_val2017_detections_AP_H_56_person.json
        │-- train2017
        │   │-- 000000000009.jpg
        │   │-- 000000000025.jpg
        │   │-- 000000000030.jpg
        │   │-- ...
        `-- val2017
            │-- 000000000139.jpg
            │-- 000000000285.jpg
            │-- 000000000632.jpg
            │-- ...

对于 MPII 数据集，请从 MPII Human Pose Dataset 下载。 我们已将原始标注文件转换为 JSON 格式，请从 mpii_annotations 下载。 将其解压到 $LITE_HRNET/data 目录下，使其呈现如下结构：

lite_hrnet
├── configs
├── models
├── tools
`── data
    │── mpii
        |── annotations
        |   |── mpii_gt_val.mat
        |   |── mpii_test.json
        |   |── mpii_train.json
        |   |── mpii_trainval.json
        |   `── mpii_val.json
        `── images
            |── 000001163.jpg
            |── 000003072.jpg

训练与测试

所有输出（日志文件和检查点）都将保存到工作目录中， 该目录由配置文件中的 work_dir 指定。

默认情况下，我们在每个 epoch 结束后都会对模型进行验证，您可以通过修改训练配置中的 interval 参数来调整验证频率：

evaluation = dict(interval=5)  # 每 5 个 epoch 验证一次模型。

根据 线性缩放规则，如果您使用不同数量的 GPU 或每 GPU 使用不同数量的视频，则需要将学习率设置为与批量大小成比例。例如，4 张 GPU 每张 2 个视频时 lr=0.01，而 16 张 GPU 每张 4 个视频时 lr=0.08。

训练

# 使用单张 GPU 训练
python tools/train.py ${CONFIG_FILE} [可选参数]

# 使用多张 GPU 训练
./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} [可选参数]

可选参数包括：

• --validate（强烈推荐）：在训练过程中每隔 k 个 epoch 进行一次评估（默认值为 5 个 epoch）。
• --work-dir ${WORK_DIR}：覆盖配置文件中指定的工作目录。
• --resume-from ${CHECKPOINT_FILE}：从之前的检查点文件恢复训练。
• --gpus ${GPU_NUM}：使用的 GPU 数量，仅适用于非分布式训练。
• --seed ${SEED}：为 Python、NumPy 和 PyTorch 的随机状态设置种子，以生成随机数。
• --deterministic：如果指定此选项，将为 CUDNN 后端设置确定性选项。
• JOB_LAUNCHER：用于分布式作业初始化的启动器。允许的选项有 none、pytorch、slurm、mpi。特别是，如果设置为 none，将以非分布式模式进行测试。
• LOCAL_RANK：本地排名 ID。如果未指定，将被设置为 0。
• --autoscale-lr：如果指定此选项，将根据 线性缩放规则自动按 GPU 数量缩放学习率。

resume-from 和 load-from 的区别： resume-from 会加载模型权重和优化器状态，并且继承指定检查点的 epoch。通常用于恢复意外中断的训练过程。 load-from 只加载模型权重，训练 epoch 将从 0 开始。通常用于微调。

示例：

在 COCO train2017 数据集上训练

./tools/dist_train.sh configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py 8

在 MPII 数据集上训练

./tools/dist_train.sh configs/top_down/lite_hrnet/mpii/litehrnet_18_mpii_256x256.py 8

测试

您可以使用以下命令测试数据集。

# 单 GPU 测试
python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} [--out ${RESULT_FILE}] [--eval ${EVAL_METRIC}] \
    [--proc_per_gpu ${NUM_PROC_PER_GPU}] [--gpu_collect] [--tmpdir ${TMPDIR}] [--average_clips ${AVG_TYPE}] \
    [--launcher ${JOB_LAUNCHER}] [--local_rank ${LOCAL_RANK}]

# 多GPU测试
./tools/dist_test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} [--out ${RESULT_FILE}] [--eval ${EVAL_METRIC}] \
    [--proc_per_gpu ${NUM_PROC_PER_GPU}] [--gpu_collect] [--tmpdir ${TMPDIR}] [--average_clips ${AVG_TYPE}] \
    [--launcher ${JOB_LAUNCHER}] [--local_rank ${LOCAL_RANK}]

可选参数：

• RESULT_FILE：输出结果的文件名。若未指定，则不会将结果保存到文件。
• EVAL_METRIC：对结果进行评估的指标。允许的值取决于数据集。
• NUM_PROC_PER_GPU：每个GPU上的进程数。若未指定，则单个GPU上仅分配一个进程。
• --gpu_collect：若指定，识别结果将通过GPU通信进行收集。否则，结果将被分别保存到不同GPU的TMPDIR中，并由rank 0的工作节点进行汇总。
• TMPDIR：用于从多个工作节点收集结果的临时目录，当未指定--gpu_collect时可用。
• AVG_TYPE：用于对测试片段进行平均的选项。若设置为prob，则在平均片段得分之前会先应用softmax；否则直接对片段得分进行平均。
• JOB_LAUNCHER：分布式任务初始化启动器。允许的选项包括none、pytorch、slurm、mpi。特别地，若设置为none，则将以非分布式模式进行测试。
• LOCAL_RANK：本地排名ID。若未指定，则默认为0。

示例：

使用8个GPU在COCO数据集上测试LiteHRNet-18，并评估mAP。

./tools/dist_test.sh configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py \
    checkpoints/SOME_CHECKPOINT.pth 8 \
    --eval mAP

获取计算复杂度

您可以使用以下命令来计算单个模型的复杂度。

python tools/summary_network.py ${CONFIG_FILE} --shape ${SHAPE}

参数：

• SHAPE：输入尺寸。

示例：

测试输入分辨率为256x256的LiteHRNet-18的复杂度。

python tools/summary_network.py configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py \
    --shape 256 256 \

致谢

感谢以下项目和团队：

• MMPose
• HRNet

引用

如果您在研究中使用了我们的代码或模型，请引用如下文献：

@inproceedings{Yulitehrnet21,
  title={Lite-HRNet: A Lightweight High-Resolution Network},
  author={Yu, Changqian and Xiao, Bin and Gao, Changxin and Yuan, Lu and Zhang, Lei and Sang, Nong and Wang, Jingdong},
  booktitle={CVPR},
  year={2021}
}

@inproceedings{SunXLW19,
  title={Deep High-Resolution Representation Learning for Human Pose Estimation},
  author={Ke Sun and Bin Xiao and Dong Liu and Jingdong Wang},
  booktitle={CVPR},
  year={2019}
}

@article{WangSCJDZLMTWLX19,
  title={Deep High-Resolution Representation Learning for Visual Recognition},
  author={Jingdong Wang and Ke Sun and Tianheng Cheng and 
          Borui Jiang and Chaorui Deng and Yang Zhao and Dong Liu and Yadong Mu and 
          Mingkui Tan and Xinggang Wang and Wenyu Liu and Bin Xiao},
  journal={TPAMI}
  year={2019}
}

Lite-HRNet 快速上手指南

Lite-HRNet 是一个轻量级的高分辨率网络，专为人体姿态估计设计。它在保持低计算量（FLOPs）和少参数量的同时，在 COCO 和 MPII 数据集上取得了优于 MobileNet、ShuffleNet 等主流轻量级网络的性能。

1. 环境准备

系统要求

• 操作系统: Linux (官方不支持 Windows)
• Python: 3.6+
• PyTorch: 1.3+
• CUDA: 9.2+ (若从源码编译 PyTorch，CUDA 9.0 也可兼容)
• 编译器: GCC 5+
• 硬件: NVIDIA GPU (开发测试基于 8 张 V100，其他显卡未完全测试)

前置依赖

请确保安装以下核心依赖库：

• mmcv-full (必须安装最新版)
• Numpy
• opencv-python (cv2)
• json_tricks
• xtcocotools

2. 安装步骤

第一步：安装 mmcv-full

注意：如果之前安装过 mmcv，请先执行 pip uninstall mmcv 卸载，避免冲突。

推荐使用预编译包安装（请将 {cu_version} 和 {torch_version} 替换为你实际的版本，例如 cu110 和 torch1.7.0）：

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html

示例 (CUDA 11.0 + PyTorch 1.7.0):

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu110/torch1.7.0/index.html

国内用户加速建议：如果上述链接下载缓慢，可尝试使用 pip 国内镜像源（如清华源），但需确保 -f 指定的索引地址不变：

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu110/torch1.7.0/index.html -i https://pypi.tuna.tsinghua.edu.cn/simple

若需从源码编译（适用于特殊环境）：

git clone https://github.com/open-mmlab/mmcv.git
cd mmcv
MMCV_WITH_OPS=1 pip install -e .
cd ..

第二步：安装其他依赖

克隆本项目代码后，进入目录安装其余要求：

pip install -r requirements.txt

3. 数据准备

建议将数据集根目录软链接到 $LITE_HRNET/data。

COCO 数据集

1. 下载 COCO 2017 Train/Val 图片及标注文件。
2. 下载复现多人姿态估计所需的人体检测结果文件 (OneDrive 下载)。
3. 目录结构应如下：

lite_hrnet
└── data
    └── coco
        ├── annotations
        │   ├── person_keypoints_train2017.json
        │   └── person_keypoints_val2017.json
        ├── person_detection_results
        │   └── COCO_val2017_detections_AP_H_56_person.json
        ├── train2017
        └── val2017

MPII 数据集

1. 下载 MPII 原始图片。
2. 下载转换好的 JSON 标注文件 (阿里云 OSS 下载)。
3. 目录结构应如下：

lite_hrnet
└── data
    └── mpii
        ├── annotations
        │   ├── mpii_gt_val.mat
        │   ├── mpii_test.json
        │   ├── mpii_train.json
        │   ├── mpii_trainval.json
        │   └── mpii_val.json
        └── images

4. 基本使用

所有输出（日志和模型权重）默认保存在配置文件中的 work_dir 指定目录。

训练模型

单卡训练：

python tools/train.py ${CONFIG_FILE} --validate

多卡分布式训练（推荐）：

./tools/dist_train.sh ${CONFIG_FILE} ${GPU_NUM} --validate

参数说明：

• ${CONFIG_FILE}: 配置文件路径，例如 configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py
• ${GPU_NUM}: 使用的 GPU 数量
• --validate: 强烈建议添加，用于在每个评估间隔自动在验证集上进行评估。

训练示例 (COCO 数据集, Lite-HRNet-18, 4 卡):

./tools/dist_train.sh configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py 4 --validate

测试模型

单卡测试：

python tools/test.py ${CONFIG_FILE} ${CHECKPOINT_FILE} --eval mAP

多卡分布式测试：

./tools/dist_test.sh ${CONFIG_FILE} ${CHECKPOINT_FILE} ${GPU_NUM} --eval mAP

测试示例：

./tools/dist_test.sh configs/top_down/lite_hrnet/coco/litehrnet_18_coco_256x192.py work_dirs/litehrnet_18/latest.pth 4 --eval mAP

提示：根据线性缩放规则 (Linear Scaling Rule)，如果你改变了 GPU 数量或每卡的 Batch Size，请按比例调整学习率 (lr)。例如：4 卡 x2 batch/gpu 对应 lr=0.01，则 16 卡 x4 batch/gpu 对应 lr=0.08。可使用 --autoscale-lr 参数自动调整。