OpenSTL:方便好用的时空预测开源库
时空预测学习是一种学习范式,它使得模型能够通过在无监督的情况下从给定的过去帧预测未来帧,从而学习空间和时间的模式。尽管近年来取得了显著的进展,但由于不同的设置、复杂的实现和难以复现性,对其缺乏系统性的理解。我们提出了OpenSTL,这是一个全面的时空预测学习基准,将常见的方法分为recurrent-based和recurrent-free的模型两大类。OpenSTL提供了一个模块化和可扩展的框架,实现了各种最先进的方法。我们对包括合成移动物体轨迹、人体动作、驾驶场景、交通流量和天气预测等不同领域的数据集进行了标准评估。我们提供了详尽的标准评估结果和可视化样例,并发布了对应的模型权重和log记录文件。
OpenSTL的特性:
- 灵活的代码设计。OpenSTL将STL算法分解为方法(训练和预测)、模型(网络架构)和模块,并提供统一的实验API。用户可以根据不同的STL任务使用灵活的训练策略和网络开发自己的STL算法。
- 标准化基准。OpenSTL将支持STL算法的标准化基准,包括训练和评估,类似于许多开源项目(例如MMDetection和USB等)。
- 支持多种模型和任务。OpenSTL包含了十四种有代表性的时空预测学习算法和二十四种模型,涵盖了从合成移动物体轨迹到真实世界的人体动作、驾驶场景、交通流量和天气预测等六类任务和十余个数据集。
论文:https://arxiv.org/pdf/2306.11249
代码:https://github.com/chengtan9907/OpenSTL
OpenSTL框架
OpenSTL是基于PyTorch开发的时空预测代码框架,包含了多种常用的算法和模型,提供了统一的训练、评估接口。此外,我们还提供了便捷的可视化功能,便于研究和应用过程中的效果展示。OpenSTL是高度模块化、可拓展的,用户可以灵活地基于OpenSTL来开发新的算法。
安装
OpenSTL提供了conda环境设置文件,用户可以通过以下命令轻松复现环境:
git clone https://github.com/chengtan9907/OpenSTL
cd OpenSTL
conda env create -f environment.yml
conda activate OpenSTL
python setup.py develop
我们提供了环境描述和数据集准备步骤,可以参考 install.md.
教程:利用自定义数据构建自己的项目
我们提供了一个使用OpenSTL在自定义数据上进行训练、评估和可视化的教程。这个教程可以帮助用户快速使用OpenSTL构建自己的项目。详细信息请参考examples/ 目录中的 tutorial.ipynb。
我们还提供了该教程的Colab演示:Colab链接.
标准化基准结果
详尽的标准化基准结果在 docs/en/model_zoos/ 中展示,可视化样例在 docs/en/visualization。
这里,我们以Moving Fashion MNIST和KittiCaltech为例,展示标准化基准评估结果。
Moving Fashion MNIST的标准评估结果如下所示:
KittiCaltech的标准评估结果如下所示:
我们以全球气象云层预测,展示可视化样例(建议在Github上看GIF动图,更清晰嗷):
最后,欢迎大家多多点star、提issue哈!