内卷神器OpenSTL开源啦 | 一个全面好用的时空预测学习基准

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

我们提供了环境描述和数据集准备步骤，可以参考(docs/en/install.md).

教程：利用自定义数据构建自己的项目

我们提供了一个使用OpenSTL在自定义数据上进行训练、评估和可视化的教程。这个教程可以帮助用户快速使用OpenSTL构建自己的项目。

详细信息请参考(examples) 目录中的(examples/tutorial.ipynb)。

我们还提供了该教程的Colab演示:(https://colab.research.google.com/drive/19uShc-1uCcySrjrRP3peXf2RUNVzCjHh?usp=sharing).

标准化基准结果

详尽的标准化基准结果在(docs/en/model_zoos) 中展示，

可视化样例在(docs/en/visualization)。

这里，我们以Moving Fashion MNIST和KittiCaltech为例，展示标准化基准评估结果。

Moving Fashion MNIST的标准评估结果如下所示：

KittiCaltech的标准评估结果如下所示：

我们以全球气象云层预测，展示可视化样例（建议在Github上看GIF动图，更清晰嗷）：