OpenSTL：时空预测学习的全面基准

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

简介

OpenSTL 是一个用于时空预测学习的全面基准，涵盖了广泛的方法和多样化的任务，从合成的移动物体轨迹到真实世界场景，如人体运动、驾驶场景、交通流量和天气预报。OpenSTL 提供了一个模块化且可扩展的框架，以其用户友好性、组织性和全面性而著称。代码库被组织成三个抽象层，即核心层、算法层和用户界面层，自下而上排列。我们支持 PyTorch Lightning 实现 OpenSTL-Lightning（推荐）以及原生 PyTorch 版本 OpenSTL。

(返回顶部)

概述

主要特性与计划

• 灵活的代码设计。 OpenSTL 将 STL 算法分解为 methods（训练和预测）、models（网络架构）和 modules，同时提供统一的实验 API。用户可以开发自己的 STL 算法，采用灵活的训练策略和针对不同 STL 任务的网络。

• 标准基准。 OpenSTL 将支持 STL 算法的标准基准，包括训练和评估，类似于许多开源项目（例如 MMDetection 和 USB）。我们正在构建训练基准，并将同步更新结果。

• 计划。 我们计划在各种 STL 应用任务中，基于 SimVP 提供多种 STL 方法和 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 演示：

(返回顶部)

模型 zoo 和数据集概览

我们支持多种时空预测方法，并在各类 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)

• 时空预测学习基准（prepare_data 或 百度网盘）。

  当前支持的数据集

  • BAIR 机器人推动物体 (CoRL'2017) [下载] [配置文件]
  • Human3.6M (TPAMI'2014) [下载] [配置文件]
  • KTH 动作数据集 (ICPR'2004) [下载] [配置文件]
  • Kitti-Caltech 行人数据集 (IJRR'2013) [下载] [配置文件]
  • Kinetics-400 (ArXiv'2017) [下载] [配置文件]
  • Moving MNIST (ICML'2015) [下载] [配置文件]
  • Moving FMNIST (ICML'2015) [下载] [配置文件]
  • TaxiBJ (AAAI'2017) [下载] [配置文件]
  • WeatherBench (ArXiv'2020) [下载] [配置文件]

(返回顶部)

可视化

我们在此展示 ConvLSTM 的可视化示例。如需更详细的信息，请参阅 visualization。

• 对于合成运动目标轨迹预测和真实世界视频预测，其他方法的可视化示例可在 visualization/video_visualization.md 中找到。BAIR 和 Kinetics 并未纳入基准测试，仅用于说明目的。

• 对于交通流量预测，其他方法的可视化示例展示在 visualization/traffic_visualization.md 中。

• 对于天气预报，其他方法的可视化示例展示在 visualization/weather_visualization.md 中。

许可证

本项目采用 Apache 2.0 许可证 发布。更多信息请参阅 LICENSE 文件。

致谢

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

引用

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

@inproceedings{tan2023openstl,
  title={OpenSTL: 时空预测学习的全面基准},
  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={神经信息处理系统会议数据集与基准赛道},
  year={2023}
}
@inproceedings{gao2022simvp,
  title={Simvp: 更简单却更好的视频预测},
  author={Gao, Zhangyang and Tan, Cheng and Wu, Lirong and Li, Stan Z},
  booktitle={IEEE/CVF 计算机视觉与模式识别会议论文集},
  pages={3170--3180},
  year={2022}
}
@article{tan2022simvpv2,
  title={SimVP：迈向简单而强大的时空预测学习},
  author={Tan, Cheng and Gao, Zhangyang and Li, Siyuan and Li, Stan Z},
  journal={arXiv 预印本 arXiv:2211.12509},
  year={2022}
}
@inproceedings{tan2023temporal,
  title={时间注意力单元：迈向高效的时空预测学习},
  author={Tan, Cheng and Gao, Zhangyang and Wu, Lirong and Xu, Yongjie and Xia, Jun and Li, Siyuan and Li, Stan Z},
  booktitle={IEEE/CVF 计算机视觉与模式识别会议论文集},
  pages={18770--18782},
  year={2023}
}

贡献与联系

如需添加新功能、寻求帮助或报告与 OpenSTL 相关的错误，请在 GitHub 上打开带有“new features”、“help wanted”或“enhancement”标签的 issue 和 pull request。如有任何问题，欢迎通过电子邮件与我们联系。

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

(返回顶部)

OpenSTL 快速上手指南

OpenSTL 是一个全面的时空预测学习（Spatio-Temporal Predictive Learning）基准库，涵盖了从合成物体运动轨迹到真实场景（如人体动作、驾驶场景、交通流、天气预报）的多种任务。它提供了模块化且可扩展的框架，支持 PyTorch Lightning 版本（推荐）和原生 PyTorch 版本。

1. 环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐) 或 macOS
• Python 版本: <= 3.10.8
• GPU: 支持 CUDA 的 NVIDIA GPU（用于加速训练）
• 包管理器: Conda (推荐)

主要依赖库包括：torch, pytorch-lightning, timm, opencv-python, xarray, netcdf4, scikit-learn 等。完整依赖列表将在安装步骤中通过配置文件自动处理。

2. 安装步骤

推荐使用 Conda 进行环境管理，以下是基于官方提供的 environment.yml 进行安装的步骤：

# 1. 克隆仓库
git clone https://github.com/chengtan9907/OpenSTL
cd OpenSTL

# 2. 创建并激活 Conda 环境
# 注意：此步骤会自动安装所有列出的依赖项
conda env create -f environment.yml
conda activate OpenSTL

# 3. 以开发模式安装 OpenSTL
python setup.py develop

提示：如果下载依赖较慢，建议在运行 conda env create 前配置国内镜像源（如清华源或中科大源）。

3. 基本使用

以下是在单个 GPU 上，使用 SimVP+gSTA 模型在 Moving MNIST 数据集上进行非分布式训练的示例。

第一步：准备数据

运行脚本下载 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

参数说明：

• -d mmnist: 指定数据集为 Moving MNIST。
• --lr 1e-3: 设置学习率。
• -c ...: 指定模型配置文件路径（此处为 SimVP 搭配 gSTA 架构）。
• --ex_name ...: 指定实验名称，用于保存日志和模型权重。

自定义数据

如果您希望使用自己的数据集进行训练、评估或可视化，项目提供了详细的 Jupyter Notebook 教程。

• 本地查看：examples/tutorial.ipynb
• 在线体验：Open In Colab

更多模型库（Model Zoo）和支持的数据集列表，请参考官方文档中的 Benchmarks 部分。