OpenSTL: Open-source Toolbox for SpatioTemporal Predictive Learning

📘Documentation | 🛠️Installation | 🚀Model Zoo | 🆕News

This repository is an open-source project for spatiotemporal predictive learning, which provides benchmarks of STL methods in various scenarios. It also contains the implementation code for the paper (noted as SimVPv2):

SimVP: Towards Simple yet Powerful Spatiotemporal Predictive learning
Cheng Tan, Zhangyang Gao, Siyuan Li, Stan Z. Li.

Introduction

SimVPv2 is the journal version of our previous conference work (SimVP: Simpler yet Better Video Prediction, CVPR 2022). It is worth noticing that the hidden Translator $h$ in SimVP can be replaced by any MetaFormer block (satisfying the macro design of token mixing and channel mixing). Therefore, we can provide benchmarks of popular ConvNets and Vision Transformers based on SimVP.

(back to top)

Overview

Major Features and Plans

• Flexable Code Design. OpenSTL decomposes STL algorithms into methods (training and prediction), models (network architectures), and modules, while providing unified experiment API. Users can develop their own STL algorithms with flexible training strategies and networks for different STL tasks.

• Standard Benchmarks. OpenSTL will support standard benchmarks of STL algorithms image with training and evaluation as many open-source projects (e.g., MMDetection and USB). We are working on training benchmarks and will update results synchronizingly.

• Plans. We plan to provide benchmarks of various STL methods and MetaFormer architectures based on SimVP in various STL application tasks, e.g., video prediction, weather prediction, traffic prediction, etc. We encourage researchers interested in STL to contribute to OpenSTL or provide valuable advice!

Code Structures

• openstl/api contains an experiment runner.
• openstl/core contains core training plugins and metrics.
• openstl/datasets contains datasets and dataloaders.
• openstl/methods/ contains training methods for various video prediction methods.
• openstl/models/ contains the main network architectures of various video prediction methods.
• openstl/modules/ contains network modules and layers.
• tools/ contains the executable python files tools/train.py and tools/test.py with possible arguments for training, validating, and testing pipelines.

News and Updates

[2023-04-19] OpenSTL v0.2.0 is released. The training loop and dataloaders are fixed.

Installation

This project has provided an environment setting file of conda, users can easily reproduce the environment by the following commands:

git clone https://github.com/chengtan9907/OpenSTL
cd OpenSTL
conda env create -f environment.yml
conda activate OpenSTL
python setup.py develop

Dependencies

• argparse
• fvcore
• numpy
• hickle
• scikit-image
• scikit-learn
• torch
• timm
• tqdm
• xarray==0.19.0

Please refer to install.md for more detailed instructions.

Getting Started

Please see get_started.md for the basic usage. Here is an example of single GPU non-distributed training SimVP+gSTA on Moving MNIST dataset.

bash tools/prepare_data/download_mmnist.sh
python tools/train.py -d mmnist --lr 1e-3 -c configs/mmnist/simvp/SimVP_gSTA.py --ex_name mmnist_simvp_gsta

(back to top)

Overview of Model Zoo and Datasets

We support various spatiotemporal prediction methods and will provide benchmarks on various STL datasets. We are working on add new methods and collecting experiment results.

• Spatiotemporal Prediction Methods.

  Currently supported methods

  • ConvLSTM (NeurIPS'2015)
  • PredRNN (NeurIPS'2017)
  • PredRNN++ (ICML'2018)
  • E3D-LSTM (ICLR'2018)
  • MIM (CVPR'2019)
  • CrevNet (ICLR'2020)
  • PhyDNet (CVPR'2020)
  • MAU (NeurIPS'2021)
  • PredRNN.V2 (TPAMI'2022)
  • SimVP (CVPR'2022)
  • SimVP.V2 (ArXiv'2022)
  • TAU (CVPR'2023)
  Currently supported MetaFormer models for SimVP

  • ViT (ICLR'2021)
  • Swin-Transformer (ICCV'2021)
  • MLP-Mixer (NeurIPS'2021)
  • ConvMixer (Openreview'2021)
  • UniFormer (ICLR'2022)
  • PoolFormer (CVPR'2022)
  • ConvNeXt (CVPR'2022)
  • VAN (ArXiv'2022)
  • IncepU (SimVP.V1) (CVPR'2022)
  • gSTA (SimVP.V2) (ArXiv'2022)
  • HorNet (NeurIPS'2022)
  • MogaNet (ArXiv'2022)

• Spatiotemporal Predictive Learning Benchmarks.

  Currently supported datasets

  • Human3.6M (TPAMI'2014) [download] [config]
  • KTH Action (ICPR'2004) [download] [config]
  • KittiCaltech Pedestrian (IJRR'2013) [download] [config]
  • Moving MNIST (ICML'2015) [download] [config]
  • Moving FMNIST (ICML'2015) [download]
  • TaxiBJ (AAAI'2017) [download] [config]
  • WeatherBench (ArXiv'2020) [download] [config]

(back to top)

License

This project is released under the Apache 2.0 license. See LICENSE for more information.

Acknowledgement

OpenSTL is an open-source project for STL algorithms created by researchers in CAIRI AI Lab. We encourage researchers interested in video and weather prediction to contribute to OpenSTL! We borrow the official implementations of ConvLSTM, PredRNN variants, E3D-LSTM, MAU, CrevNet, and PhyDNet.

Citation

If you are interested in our repository or our paper, please cite the following paper:

@article{tan2022simvp,
  title={SimVP: Towards Simple yet Powerful Spatiotemporal Predictive Learning},
  author={Tan, Cheng and Gao, Zhangyang and Li, Siyuan and Li, Stan Z},
  journal={arXiv preprint arXiv:2211.12509},
  year={2022}
}

@misc{li2023openstl,
  title={OpenSTL: Open-source Toolbox for SpatioTemporal Predictive Learning},
  author={Li, Siyuan and Tan, Cheng and Gao, Zhangyang and Li, Stan Z},
  howpublished = {\url{https://github.com/chengtan9907/OpenSTL}},
  year={2023}
}

Contribution and Contact

For adding new features, looking for helps, or reporting bugs associated with OpenSTL, please open a GitHub issue and pull request with the tag "new features", "help wanted", or "enhancement". Feel free to contact us through email if you have any questions.

• Siyuan Li ([email protected]), Westlake University & Zhejiang University
• Cheng Tan ([email protected]), Westlake University & Zhejiang University
• Zhangyang Gao ([email protected]), Westlake University & Zhejiang University

(back to top)