/Fast_Human_Pose_Estimation_Pytorch

Pytorch Code for CVPR2019 paper "Fast Human Pose Estimation" https://arxiv.org/abs/1811.05419

Primary LanguagePythonApache License 2.0Apache-2.0

Fast Human Pose Estimation Pytorch

This is an unoffical implemention for paper Fast Human Pose Estimation, Feng Zhang, Xiatian Zhu, Mao Ye. Most of code comes from pytorch implementation for stacked hourglass network pytorch-pose. In this repo, we followed Fast Pose Distillation approach proposed by Fast Human Pose Estimation to improve accuracy of a lightweight network. We first trained a deep teacher network (stacks=8, standard convolution, 88.33@Mpii pckh), and used it to teach a student network (stacks=2, depthwise convolution, 84.69%@Mpii pckh). Our experiment shows 0.7% gain from knowledge distillation.

I benchmarked the light student model hg_s2_b1_mobile_fpd and got 43fps on i7-8700K via OpenVino. Details can be found from Fast_Stacked_Hourglass_Network_OpenVino

Please check the offical implementation by fast-human-pose-estimation.pytorch

Update at Feb 2019

  • Model trained by using extra unlabeled images uploaded, hg_s2_b1_mobile_fpd_unlabeled shows 0.28% extra gain from knowledge transfered from teacher on unlabeled data.
  • The key idea is inserting unlabeled images into mpii dataset. For unlabeled samples, loss comes from difference b/w teacher and student. For labeled samples, loss is the sum of teacher-vs-student and student-vs-groundtruth.

Results

hg_s8_b1: teacher model, hg_s2_b1_mobile:student model, hg_s2_b1_mobile_kd: student model trained with FPD. hg_s2_b1_mobile_fpd_unlabeled: student model trained with FPD with extral unlabeled samples.

Model in_res featrues # of Weights Head Shoulder Elbow Wrist Hip Knee Ankle Mean GFlops Link
hg_s8_b1 256 128 25.59m 96.59 95.35 89.38 84.15 88.70 83.98 79.59 88.33 28 GoogleDrive
hg_s2_b1_mobile 256 128 2.31m 95.80 93.61 85.50 79.63 86.13 77.82 73.62 84.69 3.2 GoogleDrive
hg_s2_b1_mobile_fpd 256 128 2.31m 95.67 94.07 86.31 79.68 86.00 79.67 75.51 85.41 3.2 GoogleDrive
hg_s2_b1_mobile_fpd_unlabeled 256 128 2.31m 95.94 94.11 87.18 80.69 87.03 79.17 74.82 85.69 3.2 GoogleDrive

Installation

  1. Create a virtualenv

    virtualenv -p /usr/bin/python2.7 pose_venv

  2. Clone the repository with submodule

    git clone --recursive https://github.com/yuanyuanli85/Fast_Human_Pose_Estimation_Pytorch.git

  3. Install all dependencies in virtualenv

    source posevenv/bin/activate
pip install -r requirements.txt

  4. Create a symbolic link to the images directory of the MPII dataset:

    ln -s PATH_TO_MPII_IMAGES_DIR data/mpii/images

  5. Disable cudnn for batchnorm layer to solve bug in pytorch0.4.0

    sed -i "1194s/torch\.backends\.cudnn\.enabled/False/g" ./pose_venv/lib/python2.7/site-packages/torch/nn/functional.py

Quick Demo

  • Download pre-trained modelhg_s2_b1_mobile_fpd) and save it to somewhere, i.e checkpoint/mpii_hg_s2_b1_mobile_fpd/
  • Run demo on sample image
python tools/mpii_demo.py -a hg -s 2 -b 1 --mobile True --checkpoint checkpoint/mpii_hg_s2_b1_mobile_fpd/model_best.pth.tar --in_res 256 --device cuda
  • You will see the detected keypoints drawn on image on your screen

Training teacher network

  • In our experiments, we used stack=8 input resolution=256 as teacher network
python example/mpii.py -a hg --stacks 8 --blocks 1 --checkpoint checkpoint/hg_s8_b1/
  • Run evaluation to get val score.
python tools/mpii.py -a hg --stacks 8 --blocks 1 --checkpoint checkpoint/hg_s8_b1/preds_best.mat

Training with Knowledge Distillation

  • Download teacher model's checkpoint or you can train from scratch. In our experiments, we used hg_s8_b1 as teacher.

  • Train student network with knowledge distillation from teacher

python example/mpii_kd.py -a hg --stacks 2 --blocks 1 --checkpoint checkpoint/hg_s2_b1_mobile/ mobile=True --teacher_stack 8 --teacher_checkpoint 
checkpoint/hg_s8_b1/model_best.pth.tar

Evaluation

Run evaluation to generate mat file

python example/mpii.py -a hg --stacks 2 --blocks 1 --checkpoint checkpoint/hg_s2_b1/ --resume checkpoint/hg_s2_b1/model_best.pth.tar -e
  • --resume_checkpoint is the checkpoint want to evaluate

Run tools/eval_PCKh.py to get val score

Export pytorch checkpoint to onnx

python tools/mpii_export_to_onxx.py -a hg -s 2 -b 1 --num-classes 16 --mobile True --in_res 256  --checkpoint checkpoint/model_best.pth.tar 
--out_onnx checkpoint/model_best.onnx

Here

  • --checkpoint is the checkpoint want to export
  • --out_onnx is the exported onnx file

Reference