通用时间序列预测Transformer的统一训练

Uni2TS是一个基于PyTorch的库,用于时间序列Transformer的研究和应用。该库旨在为通用时间序列Transformer的大规模预训练提供统一解决方案。Uni2TS还为时间序列预测提供了微调、推理和评估工具。

🎉 最新动态

2024年5月: Uni2TS论文被ICML 2024接收为口头报告!
2024年3月: 发布Uni2TS库,以及Moirai-1.0-R和LOTSA数据!

✅ 待办事项

改进文档字符串和文档

⚙️ 安装

克隆仓库:

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

创建虚拟环境:

virtualenv venv
. venv/bin/activate

从源代码构建:

pip install -e '.[notebook]'

创建一个.env文件:

touch .env

🏃 入门

让我们看一个简单的例子,展示如何使用Uni2TS从预训练模型进行零样本预测。我们首先使用pandas加载数据,以宽DataFrame的形式。 Uni2TS依赖GluonTS进行推理,因为它提供了许多时间序列预测的便利功能,如将数据集分割为训练/测试集并执行滚动评估,如下所示。

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


SIZE = "small"  # 模型大小:从{'small', 'base', 'large'}中选择
PDT = 20  # 预测长度:任意正整数
CTX = 200  # 上下文长度:任意正整数
PSZ = "auto"  # 补丁大小:从{"auto", 8, 16, 32, 64, 128}中选择
BSZ = 32  # 批量大小:任意正整数
TEST = 100  # 测试集长度:任意正整数

# 将数据读入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)

# 转换为GluonTS数据集
ds = PandasDataset(dict(df))

# 分割为训练/测试集
train, test_template = split(
    ds, offset=-TEST
)  # 将最后TEST个时间步分配为测试集

# 构建滚动窗口评估
test_data = test_template.generate_instances(
    prediction_length=PDT,  # 每个预测的时间步数
    windows=TEST // PDT,  # 滚动窗口评估中的窗口数
    distance=PDT,  # 每个窗口之间的时间步数 - distance=PDT表示非重叠窗口
)

# 准备预训练模型,从huggingface hub下载模型权重
model = MoiraiForecast(
    module=MoiraiModule.from_pretrained(f"Salesforce/moirai-1.0-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,
)

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示例

查看示例文件夹以获取更多常见任务的示例,例如可视化预测、从pandas DataFrame预测等。

💻 命令行界面

我们提供了几个脚本,作为命令行界面来轻松运行微调、评估甚至预训练任务。配置通过Hydra框架管理。

微调

首先,让我们看看如何使用Uni2TS在您的自定义数据集上微调预训练模型。 Uni2TS使用Hugging Face datasets库来处理数据加载,我们首先需要将您的数据集转换为Uni2TS格式。如果您的数据集是一个简单的pandas DataFrame,我们可以使用以下脚本轻松处理您的数据集。我们将使用流行的长序列预测基准中的ETTh1数据集作为示例。对于更复杂的用例,请参阅此笔记本以获取有关如何在Uni2TS中使用自定义数据集的更深入示例。

要开始这个过程，请在.env文件中添加你想保存处理后数据集的目录路径。

echo "CUSTOM_DATA_PATH=保存路径" >> .env

运行以下脚本，将数据集处理成所需格式。对于dataset_type选项，我们支持wide、long和wide_multivariate。

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

然而，我们在微调时可能需要验证集来进行超参数调优或早停。要额外将数据集分割成训练集和验证集，我们可以使用互斥的date_offset（日期时间字符串）或offset（整数）选项，这决定了训练集的最后一个时间步。验证集将被保存为DATASET_NAME_eval。

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

最后，我们可以简单地运行微调脚本，使用适当的训练和验证数据配置文件。

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

评估

评估脚本可用于计算评估指标，如MSE、MASE、CRPS等（参见配置文件）。

给定一个测试集（参见前面关于处理数据集的部分），我们可以运行以下命令进行评估：

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

另外，我们提供了对流行数据集的访问，可以通过数据配置进行切换。例如，假设我们想再次对来自流行的长序列预测基准的ETTh1数据集进行评估。我们首先需要下载预处理的数据集并将它们放在正确的目录中，方法是设置TSLib仓库并按照说明操作。然后，将数据集目录分配给LSF_PATH环境变量：

echo "LSF_PATH=TSLIB路径/dataset" >> .env

之后，只需使用预定义的Hydra配置文件运行以下脚本：

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

预训练

现在，让我们看看如何预训练你自己的模型。我们首先准备预训练数据，下载大规模开放时间序列档案（LOTSA数据）。假设你已经创建了一个.env文件，运行以下命令。

huggingface-cli download Salesforce/lotsa_data --repo-type=dataset --local-dir 保存路径
echo "LOTSA_V1_PATH=保存路径" >> .env

然后，我们可以简单地运行以下脚本来启动预训练任务。查看相关文件了解如何进一步自定义设置。

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

👀 引用Uni2TS

如果你在研究或应用中使用Uni2TS，请使用以下BibTeX引用：

@article{woo2024unified,
  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},
  journal={arXiv preprint arXiv:2402.02592},
  year={2024}
}