PyTorch-Pose

PyTorch-Pose 是一个基于 PyTorch 的 2D 单人人体姿态估计通用流水线实现。其目标是为最流行的人体姿态数据集（例如 MPII 人体姿态、LSP 和 FLIC）提供训练/推理/评估的接口，以及带有多种数据增强选项的数据加载器。

部分用于数据准备和增强的代码源自 Stacked Hourglass 网络，感谢原作者。

更新：该仓库现已兼容 PyTorch 0.4.1/1.0！

特性

• 多线程数据加载
• 多 GPU 训练
• 日志记录器
• 训练/测试结果可视化

安装

1. PyTorch (>= 0.4.1)：请按照 PyTorch 安装说明 进行安装。请注意，代码是基于 Python2 开发的，尚未在 Python3 上进行测试。

2. 克隆包含子模块的仓库：

   git clone --recursive https://github.com/bearpaw/pytorch-pose.git
   

3. 创建指向 MPII 数据集 images 目录的符号链接：

   ln -s PATH_TO_MPII_IMAGES_DIR data/mpii/images
   

   如果要在 COCO 数据集上进行训练/测试，请参考 COCO Readme。

4. 下载标注文件：
   • (MPII) 下载 mpii_annotations.json，并将其保存到 data/mpii
   • (MSCOCO) 下载 coco_annotations_2014.json 或/及 coco_annotations_2017.json，并将其保存到 data/mscoco

使用

详细的训练流程请参阅 TRAINING.md！

测试

您可以下载我们的预训练模型（例如 2 层 Hourglass 模型）以快速开始。

在终端中运行以下命令，即可在 MPII 验证集上评估模型（训练/验证划分来自 Tompson et al. CVPR 2015)。

CUDA_VISIBLE_DEVICES=0 python example/main.py --dataset mpii -a hg --stacks 2 --blocks 1 --checkpoint checkpoint/mpii/hg_s2_b1 --resume checkpoint/mpii/hg_s2_b1/model_best.pth.tar -e -d

• -a 指定网络架构
• --resume 将从指定模型加载权重
• -e 表示仅进行评估
• -d 会可视化网络输出，也可在训练过程中使用。

结果将保存为 .mat 文件（preds_valid.mat），它是一个 2958x16x2 的矩阵，存储在 --checkpoint 指定的文件夹中。

评估 PCKh@0.5 分数

使用 MATLAB 评估

您可以使用 MATLAB 脚本 evaluation/eval_PCKh.m 来评估您的预测结果。该评估代码移植自 Tompson et al. CVPR 2015。

使用此代码训练得到的结果（PCKh@0.5 分数）如下表所示。

模型	头部	肩部	肘部	手腕	髋部	膝盖	踝部	平均
hg_s2_b1 (最后)	95.80	94.57	88.12	83.31	86.24	80.88	77.44	86.76
hg_s2_b1 (最佳)	95.87	94.68	88.27	83.64	86.29	81.20	77.70	86.95
hg_s8_b1 (最后)	96.79	95.19	90.08	85.32	87.48	84.26	80.73	88.64
hg_s8_b1 (最佳)	96.79	95.28	90.27	85.56	87.57	84.3	81.06	88.78

训练/验证曲线如下所示。

使用 Python 评估

您也可以通过运行 python evaluation/eval_PCKh.py 来评估结果，其效果与 MATLAB 完全相同。感谢 @sssruhan1 的 贡献。

训练

在终端中运行以下命令，即可在 MPII 人体姿态数据集上训练一个 8 层 Hourglass 网络。

CUDA_VISIBLE_DEVICES=0 python example/main.py --dataset mpii -a hg --stacks 8 --blocks 1 --checkpoint checkpoint/mpii/hg8 -j 4

其中，

• CUDA_VISIBLE_DEVICES=0 指定了要使用的 GPU 设备。例如，如果要使用 ID 为 0 和 1 的两个 GPU，则应设置为 CUDA_VISIBLE_DEVICES=0,1。
• -j 指定了用于数据加载的工作进程数量。
• --checkpoint 指定了保存模型、日志和预测结果的路径。

其他

支持的数据集

• MPII 人体姿态
• 利兹运动姿态 (LSP)
• MSCOCO (单人)

支持的模型

• Stacked Hourglass 网络
• Xiao 等人，人体姿态估计与跟踪的简单基线，ECCV 2018 (PDF | GitHub)

贡献

如果您想贡献代码，请创建一个 Pull Request。

PyTorch-Pose 快速上手指南

PyTorch-Pose 是一个基于 PyTorch 的 2D 单人姿态估计通用流水线实现，支持训练、推理和评估，并提供了针对 MPII、LSP 和 FLIC 等主流数据集的数据加载与增强接口。

环境准备

• 操作系统: Linux (推荐)
• Python 版本: Python 2.7 (代码基于 Python 2 开发，尚未在 Python 3 上测试)
• 深度学习框架: PyTorch >= 0.4.1 (兼容 1.0)
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (支持多卡训练)

注意：请确保已安装 CUDA 驱动及对应的 PyTorch 版本。国内用户可通过清华或阿里镜像源加速 PyTorch 安装。

安装步骤

1. 克隆仓库

使用 --recursive 参数克隆仓库以包含子模块：

git clone --recursive https://github.com/bearpaw/pytorch-pose.git
cd pytorch-pose

2. 配置数据集路径

以 MPII 数据集为例，创建图像目录的软链接（请将 PATH_TO_MPII_IMAGES_DIR 替换为实际路径）：

ln -s PATH_TO_MPII_IMAGES_DIR data/mpii/images

若需使用 COCO 数据集，请参考 data/mscoco/README.md 进行配置。

3. 下载标注文件

将以下标注文件下载到指定目录：

• MPII 数据集: 下载 mpii_annotations.json 并保存至 data/mpii/ 目录。

• MSCOCO 数据集 (可选): 下载 coco_annotations_2014.json 或 coco_annotations_2017.json 并保存至 data/mscoco/ 目录。

基本使用

快速测试 (推理与评估)

下载预训练模型（例如 2-stack hourglass 模型），解压后放置于 checkpoint/mpii/hg_s2_b1/ 目录下。

运行以下命令在 MPII 验证集上进行评估并可视化结果：

CUDA_VISIBLE_DEVICES=0 python example/main.py --dataset mpii -a hg --stacks 2 --blocks 1 --checkpoint checkpoint/mpii/hg_s2_b1 --resume checkpoint/mpii/hg_s2_b1/model_best.pth.tar -e -d

参数说明：

• -a: 指定网络架构 (如 hg 代表 Hourglass)。
• --resume: 加载指定权重的模型文件。
• -e: 仅执行评估模式。
• -d: 可视化网络输出结果。

执行完成后，预测结果将保存为 .mat 文件 (如 preds_valid.mat)。

开始训练

以下命令演示如何在单张 GPU 上训练一个 8-stack 的 Hourglass 网络：

CUDA_VISIBLE_DEVICES=0 python example/main.py --dataset mpii -a hg --stacks 8 --blocks 1 --checkpoint checkpoint/mpii/hg8 -j 4

参数说明：

• CUDA_VISIBLE_DEVICES=0: 指定使用的 GPU ID (多卡可使用 0,1)。
• -j 4: 指定数据加载的线程数。
• --checkpoint: 指定模型、日志和预测结果的保存路径。

提示：详细的训练配方和高级用法请参阅项目根目录下的 TRAINING.md 文件。