OpenSTL: A Comprehensive Benchmark of Spatio-Temporal Predictive Learning
Introduction
OpenSTL is an exhaustive benchmark dedicated to spatio-temporal predictive learning. This framework is versatile, handling a variety of tasks from synthetic moving object paths to real-world applications such as predicting human movements, driving scenes, traffic flow, and weather conditions. The flexibility and user-friendly nature of OpenSTL make it a valuable resource for exploring and implementing predictive learning methods.
OpenSTL is structured into three primary layers for ease of use and scalability:
- Core Layer: Includes core training utilities and metrics.
- Algorithm Layer: Houses the algorithms necessary for training and predictions.
- User Interface Layer: Tailored for user interaction and customization.
The project supports both PyTorch Lightning and a straightforward PyTorch version, making it accessible for different use-case scenarios.
Major Features and Plans
-
Flexible Code Design: OpenSTL’s code is modular. It separates STL algorithms into 'methods', 'models', and 'modules', facilitated by a unified experiment API, which allows users to develop their own algorithms with tailored training strategies.
-
Standard Benchmarks: The project aims to support a wide array of standard benchmarks for STL algorithms, similar to other open-source endeavors like MMDetection. These benchmarks will be crucial for training and evaluation purposes.
-
Future Plans: OpenSTL plans to support additional STL methods and architectures across various application tasks, such as video and weather forecasts. Researchers are encouraged to contribute and enhance these endeavors.
Code Structure
The codebase is organized as follows:
openstl/api: Contains the experiment runner.openstl/core: Houses core training plugins and metrics.openstl/datasets: Includes datasets and data loaders.openstl/methods/: Hosts training methods for different prediction tasks.openstl/models/: Encompasses network architectures for prediction tasks.openstl/modules/: Consists of network modules and layers.tools/: Provides executable Python files for training, validation, and testing.
News and Updates
Recent updates have included:
- The release of OpenSTL-Lightning as version 1.0.0.
- Acceptance of the OpenSTL paper at the NeurIPS 2023 Dataset and Benchmark Track.
- Ongoing enhancements in version 0.3.0 and beyond.
Installation
Setting up OpenSTL is straightforward, facilitated by a Conda environment file. Users can replicate the project's environment using a few command line instructions to clone the repository and install dependencies.
Getting Started
To get started, users are encouraged to refer to the get_started.md guide, which provides examples including single GPU training on datasets like Moving MNIST.
Custom Data Tutorial
A tutorial is available to help users implement OpenSTL on custom datasets, enabling them to build projects tailored to specific needs. This tutorial, along with a Google Colab demo, aids in understanding and visualizing the framework's capabilities.
Model Zoo and Datasets
OpenSTL supports a range of spatiotemporal prediction methods and benchmarks across diverse datasets. These include well-known architectures like ConvLSTM, PredRNN, and newer models such as SimVP. The framework provides configurations for various datasets spanning from human movement datasets to weather prediction sets.
Visualization
OpenSTL offers robust visualization examples for different prediction tasks, ranging from synthetic datasets to real-world applications in traffic and weather forecasting. These visualizations demonstrate the practical application of various prediction methods included within the framework.
License and Acknowledgement
OpenSTL is licensed under Apache 2.0, encouraging open collaboration and development. The project acknowledges contributions from researchers and implementations from several widespread models and architectures.
Citation
Researchers and developers interested in citing the OpenSTL framework or paper can refer to the provided citation format.
OpenSTL stands as a significant tool in the field of spatio-temporal predictive learning, offering a comprehensive, organized, and modular approach that welcomes community contributions and advancements.