OpenSTL:时空预测学习的综合基准

📘文档 | 🛠️安装 | 🚀模型库 | 🤗Hugging Face | 👀可视化 | 🆕新闻

简介

OpenSTL是一个时空预测学习的综合基准,涵盖了广泛的方法和多样化的任务,从合成移动物体轨迹到现实世界场景,如人体运动、驾驶场景、交通流量和天气预报。OpenSTL提供了一个模块化和可扩展的框架,在用户友好性、组织性和全面性方面表现出色。代码库分为三个抽象层,即从底层到顶层的核心层、算法层和用户界面层。我们支持PyTorch Lightning实现OpenSTL-Lightning(推荐)和朴素PyTorch版本OpenSTL。

(返回顶部)

概述

主要特点和计划

灵活的代码设计。 OpenSTL将STL算法分解为方法(训练和预测)、模型(网络架构)和模块,同时提供统一的实验API。用户可以为不同的STL任务开发自己的STL算法,具有灵活的训练策略和网络。
标准基准。 OpenSTL将支持STL算法的标准基准,包括训练和评估,就像许多开源项目一样(例如,MMDetection和USB)。我们正在进行训练基准测试,并将同步更新结果。
计划。 我们计划在各种STL应用任务中提供各种STL方法和基于SimVP的MetaFormer架构的基准,例如视频预测、天气预测、交通预测等。我们鼓励对STL感兴趣的研究人员为OpenSTL做出贡献或提供有价值的建议!

代码结构

openstl/api包含一个实验运行器。
openstl/core包含核心训练插件和指标。
openstl/datasets包含数据集和数据加载器。
openstl/methods/包含各种视频预测方法的训练方法。
openstl/models/包含各种视频预测方法的主要网络架构。
openstl/modules/包含网络模块和层。
tools/包含可执行的Python文件tools/train.py和tools/test.py,带有训练、验证和测试管道的可能参数。

新闻和更新

[2023-12-15] OpenSTL-Lightning (OpenSTL v1.0.0)发布。

[2023-09-23] OpenSTL论文已被NeurIPS 2023数据集和基准测试轨道接受! arXiv / 知乎。

[2023-06-19] OpenSTL v0.3.0发布,并将在#25中进行增强。

安装

本项目提供了conda环境设置文件,用户可以通过以下命令轻松重现环境:

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

依赖项

argparse
dask
decord
fvcore
hickle
lpips
matplotlib
netcdf4
numpy
opencv-python
packaging
pandas
python<=3.10.8
scikit-image
scikit-learn
torch
timm
tqdm
xarray==0.19.0

请参考install.md获取更详细的说明。

入门

请参阅get_started.md了解基本用法。这里是一个使用单个GPU非分布式训练SimVP+gSTA在Moving MNIST数据集上的示例。

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

使用自定义数据的教程

为了方便用户,我们提供了一个关于如何使用OpenSTL在自定义数据上进行训练、评估和可视化的教程。这个教程使用户能够快速使用OpenSTL构建自己的项目。更多详细信息,请参阅examples/目录中的tutorial.ipynb。

我们还提供了这个教程的Colab演示:

(返回顶部)

模型库和数据集概述

我们支持各种时空预测方法,并在各种STL数据集上提供基准。我们正在努力添加新方法并收集实验结果。

时空预测方法。
当前支持的方法
- ConvLSTM (NeurIPS'2015)
- PredRNN (NeurIPS'2017)
- PredRNN++ (ICML'2018)
- E3D-LSTM (ICLR'2018)
- MIM (CVPR'2019)
- PhyDNet (CVPR'2020)
- MAU (NeurIPS'2021)
- PredRNN.V2 (TPAMI'2022)
- SimVP (CVPR'2022)
- SimVP.V2 (ArXiv'2022)
- TAU (CVPR'2023)
- MMVP (ICCV'2023)
- SwinLSTM (ICCV'2023)
- WaST (AAAI'2024)
当前支持的SimVP的MetaFormer模型
- ViT (Vision Transformer) (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)
时空预测学习基准 (准备数据或百度网盘)。
当前支持的数据集
- BAIR Robot Pushing (CoRL'2017) [下载] [配置]
- Human3.6M (TPAMI'2014) [下载] [配置]
- KTH Action (ICPR'2004) [下载] [配置]
- KittiCaltech Pedestrian (IJRR'2013) [下载] [配置]
- Kinetics-400 (ArXiv'2017) [下载] [配置]
- Moving MNIST (ICML'2015) [下载] [配置]
- Moving FMNIST (ICML'2015) [[下载](https://pan.baidu.com/s/1fudsBHyrf3nbt-7d42 | 移动 MNIST-CIFAR | KittiCaltech | | :---: | :---: | |
  |
  |

KTH	Human 3.6M

交通 - 流入	交通 - 流出

天气 - 温度	天气 - 湿度

天气 - 纬度风	天气 - 云量

BAIR 机器人推动	Kinetics-400

许可证

本项目在 Apache 2.0 许可证下发布。更多信息请参见 LICENSE。

致谢

OpenSTL 是由 CAIRI AI 实验室的研究人员创建的用于 STL 算法的开源项目。我们鼓励对视频和天气预测感兴趣的研究人员为 OpenSTL 做出贡献！我们借鉴了 ConvLSTM、PredNet、PredRNN 变体、E3D-LSTM、MAU、PhyDNet、MMVP 和 SwinLSTM 的官方实现。

引用

如果您对我们的仓库或论文感兴趣，请引用以下论文：

@inproceedings{tan2023openstl,
  title={OpenSTL: A Comprehensive Benchmark of Spatio-Temporal Predictive Learning},
  author={Tan, Cheng and Li, Siyuan and Gao, Zhangyang and Guan, Wenfei and Wang, Zedong and Liu, Zicheng and Wu, Lirong and Li, Stan Z},
  booktitle={Conference on Neural Information Processing Systems Datasets and Benchmarks Track},
  year={2023}
}

@inproceedings{gao2022simvp,
  title={Simvp: Simpler yet better video prediction},
  author={Gao, Zhangyang and Tan, Cheng and Wu, Lirong and Li, Stan Z},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={3170--3180},
  year={2022}
}
@article{tan2022simvpv2,
  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}
}
@inproceedings{tan2023temporal,
  title={Temporal attention unit: Towards efficient spatiotemporal predictive learning},
  author={Tan, Cheng and Gao, Zhangyang and Wu, Lirong and Xu, Yongjie and Xia, Jun and Li, Siyuan and Li, Stan Z},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={18770--18782},
  year={2023}
}

贡献和联系

对于添加新功能、寻求帮助或报告与 OpenSTL 相关的错误，请使用"新功能"、"需要帮助"或"改进"标签打开 GitHub issue 和 pull request。如果您有任何问题，请随时通过电子邮件与我们联系。

李思远 (lisiyuan@westlake.edu.cn)，西湖大学 & 浙江大学
谭程 (tancheng@westlake.edu.cn)，西湖大学 & 浙江大学
高张扬 (gaozhangyang@westlake.edu.cn)，西湖大学 & 浙江大学

(返回顶部)