This repository contains a custom implementation of Temporal Fusion Transformers for Interpretable Multi-horizon Time Series Forecasting using Lightning and pytorch_forecasting for demand forecasting on the Stallion dataset.
Enhancements compared to the original implementation in the Google Research repo:
- capabilities added through pytorch_forecasting base model e.g. monotone constraints
- static variables can be continuous
- multiple categorical variables can be summarized with an EmbeddingBag
- variable encoder and decoder length by sample
- categorical embeddings are not transformed by variable selection network (because it is a redundant operation)
- variable dimension in variable selection network are scaled up via linear interpolation to reduce number of parameters
- non-linear variable processing in variable selection network can be shared among decoder and encoder (not shared by default)
The dependency management system is poetry. Install poetry and run:
poetry install
to set up the environment.
You can get the baseline performance in the dataset by training pytorch_forecasting's Baseline model:
cd src
poetry run python baseline.py
To train the full model and get the performance:
poetry run python train.py
poetry run python evaluate.py
To run prediction on test data in the model, run:
poetry run python predict.py
To inference the model through an API, the api.py
file sets up a simple Flask API with a '/predictions' GET endpoint.
poetry run python api.py
then you can access localhost:8501/docs
to see a nice UI set up by FastAPI that allows easy inference, or use curl
or any other tools to call the API directly.