Uni2TS: 统一的通用时间序列预测Transformer模型

uni2ts

Uni2TS:统一的通用时间序列预测Transformer模型

时间序列预测在各行各业的决策过程中扮演着至关重要的角色。然而,与自然语言处理和图像识别等领域相比,时间序列预测中高级人工智能技术的应用相对滞后。虽然基础人工智能在自然语言处理和图像识别等领域取得了重大进展,但直到最近其对时间序列预测的影响还比较有限。不过,专门针对时间序列预测的基础模型开发正在加速发展。本文将探讨时间序列预测基础AI模型的最新进展,重点介绍Salesforce AI Research开发的Uni2TS模型。

Uni2TS简介

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

Uni2TS架构图

Uni2TS的主要特性

大规模预训练: Uni2TS使用大规模开放时间序列数据集(LOTSA)进行预训练,该数据集包含来自9个不同领域的270亿个观测值。
多尺度时间模式捕捉: 通过使用多个补丁大小投影层,Uni2TS能够捕捉不同频率的时间模式。
任意变量注意力机制: 允许跨任意数量的变量进行预测。
零样本和少样本预测: 预训练模型可以直接用于零样本预测,也可以通过微调进行少样本学习。
易用性: 提供简单的API和命令行界面,方便用户进行微调、评估和预训练。

安装和使用

Uni2TS的安装非常简单,只需几个步骤即可完成:

克隆仓库:

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

创建虚拟环境:

virtualenv venv
. venv/bin/activate

从源代码构建:

pip install -e '.[notebook]'

创建.env文件

快速开始

让我们看一个简单的例子,展示如何使用Uni2TS从预训练模型进行零样本预测:

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  # 测试集长度: 任意正整数

# 读取数据
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_data = test_template.generate_instances(
    prediction_length=PDT,
    windows=TEST // PDT,
    distance=PDT,
)

# 准备预训练模型
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()

这个例子展示了如何使用Uni2TS加载预训练模型、处理数据、进行预测并可视化结果。

高级功能

Uni2TS还提供了一系列高级功能,包括:

微调: 可以在自定义数据集上微调预训练模型。
评估: 提供多种评估指标,如MSE、MASE、CRPS等。
预训练: 支持用户使用自己的数据集进行模型预训练。

这些功能都可以通过命令行界面轻松实现,具体使用方法可以参考Uni2TS的官方文档。

总结

Uni2TS作为一个统一的通用时间序列预测Transformer模型,为时间序列预测领域带来了新的可能性。它不仅提供了强大的预训练模型,还为研究人员和开发者提供了灵活的工具,以便进行进一步的探索和应用。随着时间序列数据在各个领域的重要性日益凸显,Uni2TS有望成为推动这一领域发展的重要工具。

如果您在研究或应用中使用Uni2TS,欢迎引用相关论文:

@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}
}

未来,随着更多研究者和开发者的参与,Uni2TS有望进一步完善和扩展其功能,为时间序列预测领域带来更多创新和突破。我们期待看到Uni2TS在各个领域的应用,以及它如何推动时间序列分析和预测技术的进步。

Uni2TS预测结果示例