/pytorch_stacked_hourglass_cutout

Improving Stacked Hourglass Networks with Regularization

Primary LanguagePythonBSD 3-Clause "New" or "Revised" LicenseBSD-3-Clause

Improving Stacked Hourglass Networks with Regularization

Extension of Stacked Hourglass Networks for Human Pose Estimation. Alejandro Newell, Kaiyu Yang, and Jia Deng. European Conference on Computer Vision (ECCV), 2016. Github

Overview

Newell et al. originally reported 0.881 validation accuracy using 8HG model on MPII. Here we get validation accuracy of 0.885 using a 2HG model and 0.901 using an 8HG model. In this implementation, validation accuracies of 0.887 and 0.906 are achieved by adding mean-normalization, cutout, and vertical flipping. Test number of 0.913 is also achieved, as opposed to authors' 0.909.

Getting Started

This repository provides everything necessary to train and evaluate a single-person pose estimation model on MPII. If you plan on training your own model from scratch, we highly recommend using multiple GPUs.

Requirements:

  • Python 3 (code has been tested on Python 3.6)
  • PyTorch (code tested with 1.0), torchvision (tested with version 0.2.1)
  • CUDA and cuDNN
  • Python packages (not exhaustive): opencv-python, tqdm, cffi, h5py, scipy (tested with 1.1.0)

Structure:

  • data/: data loading and data augmentation code
  • models/: network architecture definitions
  • task/: task-specific functions and training configuration
  • utils/: image processing code and miscellaneous helper functions
  • train.py: code for model training
  • test.py: code for model evaluation

Dataset

Download the full MPII Human Pose dataset, and place the images directory in data/MPII/

Training and Testing

To train a network, call:

python train.py -e test_run_001 (-e,--exp allows you to specify an experiment name)

To continue an experiment where it left off, you can call:

python train.py -c test_run_001

All training hyperparameters are defined in task/pose.py, and you can modify __config__ to test different options. It is likely you will have to change the batchsize to accommodate the number of GPUs you have available.

Once a model has been trained, you can evaluate it with:

python test.py -c test_run_001

The option "-m n" will automatically stop training after n total iterations (if continuing, would look at total iterations)

Pretrained Models

An 8HG pretrained model is available here. It should yield validation accuracy of 0.906.

A 2HG pretrained model is available here. It should yield validation accuracy of 0.887.

Models should be formatted as exp/<exp_name>/checkpoint.pt

Training/Validation split

The train/val split is same as that found in authors' implementation

Note

During training, occasionaly "ConnectionResetError" warning was displayed between epochs, but did not affect training.

Special Thanks

PyTorch code extended from here. Implemented for a project under advisors Alejandro Newell and Prof. Jia Deng.