DOCSEARCHbyAlgolia
Submit Docs

uni2ts: Unified Training of Universal Time Series Forecasting Transformers

repository·main·Indexed 23 days ago

https://github.com/salesforceairesearch/uni2ts

A PyTorch-based library for large-scale pre-training, fine-tuning, inference, and evaluation of Universal Time Series Transformers, including the Moirai models. It supports benchmarking for Chronos, TimesFM, and VisionTS, and provides tools for reproducing Moirai-1.0-R pre-training and evaluation across Monash TSF, Probabilistic Forecasting (PF), and Long Sequence Forecasting (LSF) datasets. The repository also includes the Moirai Agent, a framework for context-aware forecasting and expert selection using LLMs.

Tokens
39.2K
Snippets
79
Records
146
Agent score
82%

What's inside uni2ts

  1. Overview of Moirai-MoE-1.0-R

    main
    Moirai-MoE is a Sparse Mixture-of-Experts (MoE) time series foundation model. Unlike the original Moirai which uses multi-heuristic-defined projection layers for different frequencies, Moirai-MoE uses a single input/output projection layer and delegates pattern capture to sparse MoE Transformers. This allows for data-driven specialization at the token level.

  2. Overview of Moirai Agent

    main

    Moirai Agent is an intelligent forecasting framework designed to blend historical time series data with external contextual signals. Unlike traditional numerical-only models, it uses Large Language Models (LLMs) as orchestrators to integrate heterogeneous information, select forecasting experts, and adjust predictions based on real-world context.

    It provides two core capabilities:

    1. Context-aware forecasting: Combines numerical history with external contextual inputs to refine lookback windows, detect anomalies, and anticipate future effects.
    2. Expert selection: Uses an LLM-based selector to choose the most suitable forecasting model for a specific task from a pool of state-of-the-art forecasters.

  3. Reproduce Moirai-1.0-R Pre-training

    main

    To reproduce the pre-training runs, execute the cli.train module from the base uni2ts folder. You may need to adjust gradient accumulation settings depending on your hardware.

    Note that the data parameter should be set to lotsa_v1_weighted to use the Large-scale Open Time Series Archive (LOTSA) dataset.

    # Example: Pre-training the Small model
python -m cli.train \ 
  -cp conf/pretrain \ 
  run_name=moirai_small \ 
  model=moirai_small \ 
  data=lotsa_v1_weighted \ 
  trainer.max_epochs=1_000 \ 
  train_dataloader.batch_size=256

  6. Reproduce the VN1 Forecasting Challenge results

    main

    To reproduce the first-place results achieved by the fine-tuned Moirai-base model in the VN1 Forecasting challenge, follow these steps:

    1. Environment Setup: Follow the standard uni2ts library instructions to create a virtual environment and install all necessary dependencies.
    2. Download Raw Data: Use the provided Makefile to download the required raw dataset:
      make download_data
    3. Preprocess Data: Run prepare_data.py after updating the script with the correct directory path of your downloaded raw dataset.
    4. Configure Environment: Set the CUSTOM_DATA_PATH environment variable in your .env file to point to your processed dataset directory:
      echo "CUSTOM_DATA_PATH=PATH_TO_SAVE" >> .env
    5. Fine-tune Moirai-base:
      • Open the configuration file and replace pretrained_model_name_or_path with your local path to the pretrained model.
      • Execute the training command:
      python -m cli.train -cp ../project/vn1_competition/fine_tune run_name=run1
    6. Run Inference: Open src/main.py, update the weight file path to point to your fine-tuned model weights, and execute the script.
    make download_data

echo "CUSTOM_DATA_PATH=PATH_TO_SAVE" >> .env

python -m cli.train -cp ../project/vn1_competition/fine_tune run_name=run1

  7. Create LSF datasets for fine-tuning

    main

    To fine-tune Moirai-1.0-R, you must create offline datasets using sliding windows of fixed-length time series samples.

    Key Configuration Details:

    • Dataset Type: Set dataset_type to "wide_multivariate" for multivariate setups or "wide" for univariate setups.
    • Sliding Window Distance: The number of training samples is determined by the sliding window distance. The default is distance=1. For large datasets, increasing this value reduces computational cost per epoch.
    • Normalization: The LSF setup requires normalizing data using training statistics.
    • Split Configuration: Dataset splits follow the logic in src/uni2ts/eval_util/_lsf_dataset.py. Configuration files for these splits are located under cli/conf/finetune/data and cli/conf/finetune/val_data.

    Implementation Files:

    • Dataset creation logic is found in src/uni2ts/data/builder/simple.py.
    • Creation scripts can be found in project/moirai-1/finetune_lsf/build_lsf_ft_datasets.sh.

  9. Evaluate using popular LSF datasets

    main

    You can evaluate models on standard Long Sequence Forecasting (LSF) benchmarks by using predefined Hydra configurations.

    1. Download the required datasets via the TSLib repository.
    2. Set the LSF_PATH environment variable in your .env file to point to the dataset directory.
    3. Run cli.eval using the lsf_test data configuration and specify the dataset_name and prediction_length.
    # 1. Set LSF path
echo "LSF_PATH=PATH_TO_TSLIB/dataset" >> .env

# 2. Run evaluation
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

  10. Explore Moirai Agent workflows

    main

    The Moirai Agent repository is organized into two primary functional areas for research and replication:

    • ctx_forecast: Contains detailed scripts for implementing and running context-aware forecasting workflows.
    • gift_eval: Contains scripts focused on replicating gift_eval results using standard numeric forecasting through model selection.

    For benchmarking information, standard forecasting results are evaluated on the GIFT-Eval leaderboard, while contextual forecasting results are evaluated on the GIFT-CTX Dataset.