Unified Training of Universal Time Series Transformers

Uni2TS is a PyTorch based library for research and applications related to Time Series Forecasting. It provides a unified framework for large-scale pre-training, fine-tuning, inference, and evaluation of Universal Time Series Transformers.

Related reading: Moirai Paper, Moirai Salesforce Blog, Moirai-MoE Paper, Moirai-MoE Salesforce Blog, Moirai-MoE AI Horizon Forecast Blog, Moirai-MoE Jiqizhixin Blog.

🎉 What's New

• Nov 2024: The first general time series forecasting benchmark GIFT-Eval is released. Moirai-Large achieves the best performance in the Leaderboard!

• Oct 2024: A new model Moirai-MoE! The preprint is available on arXiv, along with the model weights of Moirai-MoE-Small and Moirai-MoE-Base. Getting started with inference code and notebook examples!

• Sep 2024: Released Evaluation Code of TimesFM, Chronos and VisionTS on Monash, LSF and PF benchmarks.

• Jun 2024: Released Moirai-1.1-R model weights in small, base, and large.

• May 2024: The Moirai Paper has been accepted to ICML 2024 as an Oral presentation!

• Mar 2024: Release of Uni2TS library, along with Moirai Paper, Moirai-1.0-R Models, and LOTSA Data.

⚙️ Installation

1. Clone repository:

git clone https://github.com/SalesforceAIResearch/uni2ts.git
cd uni2ts

2. Create virtual environment:

virtualenv venv
. venv/bin/activate

3. Build from source:

pip install -e '.[notebook]'

4. Create a .env file:

touch .env

We also support installation via PyPI.

pip install uni2ts

🏃 Getting Started

Let's see a simple example on how to use Uni2TS to make zero-shot forecasts from a pre-trained model. We first load our data using pandas, in the form of a wide DataFrame. Uni2TS relies on GluonTS for inference as it provides many convenience functions for time series forecasting, such as splitting a dataset into a train/test split and performing rolling evaluations, as demonstrated below.

import torch
import matplotlib.pyplot as plt
import pandas as pd
from gluonts.dataset.pandas import PandasDataset
from gluonts.dataset.split import split
from huggingface_hub import hf_hub_download

from uni2ts.eval_util.plot import plot_single
from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
from uni2ts.model.moirai_moe import MoiraiMoEForecast, MoiraiMoEModule

MODEL = "moirai-moe"  # model name: choose from {'moirai', 'moirai-moe'}
SIZE = "small"  # model size: choose from {'small', 'base', 'large'}
PDT = 20  # prediction length: any positive integer
CTX = 200  # context length: any positive integer
PSZ = "auto"  # patch size: choose from {"auto", 8, 16, 32, 64, 128}
BSZ = 32  # batch size: any positive integer
TEST = 100  # test set length: any positive integer

# Read data into pandas DataFrame
url = (
    "https://gist.githubusercontent.com/rsnirwan/c8c8654a98350fadd229b00167174ec4"
    "/raw/a42101c7786d4bc7695228a0f2c8cea41340e18f/ts_wide.csv"
)
df = pd.read_csv(url, index_col=0, parse_dates=True)

# Convert into GluonTS dataset
ds = PandasDataset(dict(df))

# Split into train/test set
train, test_template = split(
    ds, offset=-TEST
)  # assign last TEST time steps as test set

# Construct rolling window evaluation
test_data = test_template.generate_instances(
    prediction_length=PDT,  # number of time steps for each prediction
    windows=TEST // PDT,  # number of windows in rolling window evaluation
    distance=PDT,  # number of time steps between each window - distance=PDT for non-overlapping windows
)

# Prepare pre-trained model by downloading model weights from huggingface hub
if MODEL == "moirai":
    model = MoiraiForecast(
        module=MoiraiModule.from_pretrained(f"Salesforce/moirai-1.1-R-{SIZE}"),
        prediction_length=PDT,
        context_length=CTX,
        patch_size=PSZ,
        num_samples=100,
        target_dim=1,
        feat_dynamic_real_dim=ds.num_feat_dynamic_real,
        past_feat_dynamic_real_dim=ds.num_past_feat_dynamic_real,
    )
elif MODEL == "moirai-moe":
    model = MoiraiMoEForecast(
        module=MoiraiMoEModule.from_pretrained(f"Salesforce/moirai-moe-1.0-R-{SIZE}"),
        prediction_length=PDT,
        context_length=CTX,
        patch_size=16,
        num_samples=100,
        target_dim=1,
        feat_dynamic_real_dim=ds.num_feat_dynamic_real,
        past_feat_dynamic_real_dim=ds.num_past_feat_dynamic_real,
    )

predictor = model.create_predictor(batch_size=BSZ)
forecasts = predictor.predict(test_data.input)

input_it = iter(test_data.input)
label_it = iter(test_data.label)
forecast_it = iter(forecasts)

inp = next(input_it)
label = next(label_it)
forecast = next(forecast_it)

plot_single(
    inp, 
    label, 
    forecast, 
    context_length=200,
    name="pred",
    show_label=True,
)
plt.show()

📔 Jupyter Notebook Examples

See the example folder for more examples on common tasks, e.g. visualizing forecasts, predicting from pandas DataFrame, etc.

💻 Command Line Interface

We provide several scripts which act as a command line interface to easily run fine-tuning, evaluation, and even pre-training jobs. Configurations are managed with the Hydra framework.

Fine-tuning

Firstly, let's see how to use Uni2TS to fine-tune a pre-trained model on your custom dataset. Uni2TS uses the Hugging Face datasets library to handle data loading, and we first need to convert your dataset into the Uni2TS format. If your dataset is a simple pandas DataFrame, we can easily process your dataset with the following script. We'll use the ETTh1 dataset from the popular Long Sequence Forecasting benchmark for this example. For more complex use cases, see this notebook for more in-depth examples on how to use your custom dataset with Uni2TS.

1. To begin the process, add the path to the directory where you want to save the processed dataset into the .env file.

echo "CUSTOM_DATA_PATH=PATH_TO_SAVE" >> .env

2. Run the following script to process the dataset into the required format. For the dataset_type option, we support wide, long and wide_multivariate.

python -m uni2ts.data.builder.simple ETTh1 dataset/ETT-small/ETTh1.csv --dataset_type wide

However, we may want validation set during fine-tuning to perform hyperparameter tuning or early stopping. To additionally split the dataset into a train and validation split we can use the mutually exclusive date_offset (datetime string) or offset (integer) options which determines the last time step of the train set. The validation set will be saved as DATASET_NAME_eval.

python -m uni2ts.data.builder.simple ETTh1 dataset/ETT-small/ETTh1.csv --date_offset '2017-10-23 23:00:00'

3. Finally, we can simply run the fine-tuning script with the appropriate training and validation data configuration files.

python -m cli.train \
  -cp conf/finetune \
  run_name=example_run \ 
  model=moirai_1.0_R_small \ 
  data=etth1 \ 
  val_data=etth1  

Evaluation

The evaluation script can be used to calculate evaluation metrics such as MSE, MASE, CRPS, and so on (see the configuration file).

Given a test split (see previous section on processing datasets), we can run the following command to evaluate it:

python -m cli.eval \ 
  run_name=example_eval_1 \
  model=moirai_1.0_R_small \
  model.patch_size=32 \ 
  model.context_length=1000 \
  data=etth1_test

Alternatively, we provide access to popular datasets, and can be toggled via the data configurations. As an example, say we want to perform evaluation, again on the ETTh1 dataset from the popular Long Sequence Forecasting benchmark. We first need to download the pre-processed datasets and put them in the correct directory, by setting up the TSLib repository and following the instructions. Then, assign the dataset directory to the LSF_PATH environment variable:

echo "LSF_PATH=PATH_TO_TSLIB/dataset" >> .env

Thereafter, simply run the following script with the predefined Hydra config file:

python -m cli.eval \ 
  run_name=example_eval_2 \
  model=moirai_1.0_R_small \
  model.patch_size=32 \ 
  model.context_length=1000 \ 
  data=lsf_test \
  data.dataset_name=ETTh1 \
  data.prediction_length=96 

Pre-training

Now, let's see how you can pre-train your own model. We'll start with preparing the data for pre-training first, by downloading the Large-scale Open Time Series Archive (LOTSA data). Assuming you've already createed a .env file, run the following commands.

huggingface-cli download Salesforce/lotsa_data --repo-type=dataset --local-dir PATH_TO_SAVE
echo "LOTSA_V1_PATH=PATH_TO_SAVE" >> .env

Then, we can simply run the following script to start a pre-training job. See the relevant files on how to further customize the settings.

python -m cli.train \
  -cp conf/pretrain \
  run_name=first_run \
  model=moirai_small \
  data=lotsa_v1_unweighted

👀 Citation

If you're using this repository in your research or applications, please cite using the following BibTeX:

@article{liu2024moiraimoe,
  title={Moirai-MoE: Empowering Time Series Foundation Models with Sparse Mixture of Experts},
  author={Liu, Xu and Liu, Juncheng and Woo, Gerald and Aksu, Taha and Liang, Yuxuan and Zimmermann, Roger and Liu, Chenghao and Savarese, Silvio and Xiong, Caiming and Sahoo, Doyen},
  journal={arXiv preprint arXiv:2410.10469},
  year={2024}
}

@article{aksu2024gifteval,
  title={GIFT-Eval: A Benchmark For General Time Series Forecasting Model Evaluation},
  author={Aksu, Taha and Woo, Gerald and Liu, Juncheng and Liu, Xu and Liu, Chenghao and Savarese, Silvio and Xiong, Caiming and Sahoo, Doyen},
  journal={arXiv preprint arXiv:2410.10393},
  year={2024}
}

@inproceedings{woo2024moirai,
  title={Unified Training of Universal Time Series Forecasting Transformers},
  author={Woo, Gerald and Liu, Chenghao and Kumar, Akshat and Xiong, Caiming and Savarese, Silvio and Sahoo, Doyen},
  booktitle={Forty-first International Conference on Machine Learning},
  year={2024}
}