PoseMamba: Monocular 3D Human Pose Estimation with Bidirectional Spatio-Temporal State Space Model

Official PyTorch implementation · AAAI 2025 · Linear-complexity Mamba for 3D human pose estimation

Having trouble installing or reproducing results? See the pinned Installation & FAQ discussion before opening an issue.

---

Highlights

• First bidirectional global-local spatio-temporal SSM/Mamba design for monocular 3D HPE (2D-to-3D lifting).
• SOTA on Human3.6M & MPI-INF-3DHP with fewer parameters and MACs than Transformer baselines.
• PoseMamba-L: 38.1 mm MPJPE on H36M (P1, detected 2D), only 6.7M params / 27.9G MACs.
• 2.8× faster than MotionAGFormer; 64.7% less GPU memory at 243-frame batch inference (see paper).
• In-the-wild demo with 2D detector + 3D lifting pipeline.

Paper: AAAI 2025 · arXiv · PDF

Results on Human3.6M (243 frames, MPJPE ↓)

Method	Params	MACs	P1 (2D det)	P1 (GT 2D)
MixSTE	33.6M	139.0G	40.9	21.6
MotionBERT	42.3M	174.8G	39.2	17.8
MotionAGFormer-L	7.9M	33.0G	38.4	17.3
PoseMamba-L (ours)	6.7M	27.9G	38.1	15.6
PoseMamba-B	3.4M	13.9G	40.8	16.8
PoseMamba-S	0.9M	3.6G	41.8	20.0

Quick Start

1. Environment

Tested with Python 3.8.5, PyTorch 1.13.1+cu117, CUDA 11.7.

conda create -n posemamba python=3.8.5
conda activate posemamba
pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
pip install -r requirements.txt
cd kernels/selective_scan && pip install -e . && cd ../..

Important: kernels/selective_scan must compile successfully. If you see selective_scan_cuda_core is not defined, rebuild the CUDA extension (see FAQ).

2. Demo (in-the-wild video)

1. Download YOLOv3 + HRNet weights → ./demo/lib/checkpoint/ (Google Drive)
2. Download demo 3D lifting checkpoint → ./checkpoint/ (link)
3. Optional: download sample assets (link) and unzip
4. Run:

python vis.py --video sample_video.mp4 --gpu 0

Or use demo.sh after placing your video under ./demo/video/.

Try online: Hugging Face Demo Space · Colab notebook

Dataset

Human3.6M

1. Download MotionBERT preprocessed H3.6M data (OneDrive or Google Drive) → data/motion3d/
2. Slice clips:

cd tools && python convert_h36m.py && cd ..

Expected layout: data/motion3d/MB3D_f243s81/ with h36m_sh_conf_cam_source_final.pkl (see config data_root).

MPI-INF-3DHP

Follow MotionAGFormer for dataset setup.

Training

Config files are under configs/pose3d/ (not configs/h36m).

# Example: PoseMamba-S on Human3.6M
CUDA_VISIBLE_DEVICES=0 python train.py \
  --config configs/pose3d/PoseMamba_train_h36m_S.yaml \
  --checkpoint checkpoint/pose3d/PoseMamba_train_h36m_S

See train.sh for S/B/L variants.

Evaluation

Option A: Released .pth.tr weights (Google Drive)

Download from Hugging Face (recommended) or Google Drive:

Model	Params	Hugging Face	Google Drive
PoseMamba-S	0.9M	HF	link
PoseMamba-B	3.4M	HF	link
PoseMamba-L	6.7M	HF	link

All models: nankingjings/PoseMamba-weights · Google Drive bundle

Place weights under checkpoint/ and run evaluation following the training checkpoint layout in eval.sh.

Option B: Evaluate a training run (.bin checkpoint)

After training, checkpoints are saved as best_epoch.bin / latest_epoch.bin:

CUDA_VISIBLE_DEVICES=0 python train.py \
  --config checkpoint/pose3d/PoseMamba_B/config.yaml \
  --evaluate checkpoint/pose3d/PoseMamba_B/best_epoch.bin \
  --checkpoint eval/checkpoint

Note: training saves .bin files; released Google Drive weights may use .pth.tr. Use the format that matches your download.

FAQ

selective_scan_cuda_core is not defined

Rebuild the CUDA kernel with the same PyTorch/CUDA version as your runtime:

cd kernels/selective_scan
pip uninstall selective-scan -y 2>/dev/null || true
pip install -e .
cd ../..

Ensure nvcc matches CUDA 11.7 if using the default PyTorch build.

Dataset not found / 训练时数据集找不到

Check data_root in your yaml (e.g. data/motion3d/MB3D_f243s81/). Run tools/convert_h36m.py after downloading preprocessed data. The folder must contain the pickle file named in the config (dt_file).

Evaluation command / checkpoint format confusion

• Training output: --evaluate path/to/best_epoch.bin (see eval.sh)
• Released weights: may be .pth.tr from Google Drive; load according to your local checkpoint layout
• There is no --eval-only flag in the current train.py

Results gap vs paper

Common causes: wrong data split, not using flip test (flip: True in config), mismatched 2D detections, or selective_scan not built correctly. Please open an issue with your config, checkpoint, and log output.

How are MACs computed?

MACs follow the protocol in our paper (per-frame multiply-accumulate operations), comparable to MotionBERT / MotionAGFormer reporting.

Project structure

PoseMamba/
├── configs/pose3d/     # Training configs (S / B / L)
├── lib/model/          # PoseMamba backbone
├── kernels/selective_scan/  # CUDA SSM kernel (required)
├── tools/              # Dataset preprocessing
├── vis.py              # In-the-wild demo
└── train.py            # Train & evaluate

Acknowledgement

Code builds on MotionBERT, MotionAGFormer, P-STMO, MHFormer, and VMamba. Thank you for open-sourcing.

Citation

If you find PoseMamba useful, please cite:

@inproceedings{huang2025posemamba,
  title={PoseMamba: Monocular 3D Human Pose Estimation with Bidirectional Global-Local Spatio-Temporal State Space Model},
  author={Huang, Yunlong and Liu, Junshuo and Xian, Ke and Qiu, Robert Caiming},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  year={2025}
}

Contact

Issues & questions: Installation & FAQ · GitHub Issues · Author: Yunlong Huang