Uni2TS 是一个基于 PyTorch 的库，用于时间序列预测相关的研究和应用。它提供了一个统一的框架，用于大规模预训练、微调、推理和评估通用时间序列 Transformer。

相关阅读：Moirai 论文、Moirai Salesforce 博客、Moirai-MoE 论文、Moirai-MoE Salesforce 博客、Moirai-MoE AI Horizon Forecast 博客、Moirai-MoE Jiqizhixin 博客。

🎉 最新动态

• 2025年8月：发布了 Moirai-2.0-R-small。

• 2025年8月：增强了微调模块，并添加了在 LSF 基准上运行的 示例。

• 2024年11月：首个通用时间序列预测基准 GIFT-Eval 发布。排行榜 已上线，请尝试将您的模型提交到该平台上！

• 2024年10月：全新模型 Moirai-MoE！预印本已在 arXiv 上发布，同时提供了 Moirai-MoE-Small 和 Moirai-MoE-Base 的模型权重。可通过 推理代码 和 笔记本示例 开始使用！

• 2024年9月：在 Monash、LSF 和 PF 基准上发布了 TimesFM、Chronos 和 VisionTS 的 评估代码。

• 2024年6月：发布了 Moirai-1.1-R 模型权重，包括 small、base 和 large 版本。

• 2024年5月：Moirai 论文 被 ICML 2024 接受为口头报告！

• 2024年3月：Uni2TS 库正式发布，同时发布了 Moirai 论文、Moirai-1.0-R 模型 和 LOTSA 数据集。

⚙️ 安装

1. 克隆仓库：

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

2. 创建虚拟环境：

virtualenv venv
. venv/bin/activate

3. 从源码构建：

pip install -e '.[notebook]'

4. 创建 .env 文件：

touch .env

我们还支持通过 PyPI 安装。

pip install uni2ts

🏃 快速入门

让我们来看一个简单的例子，演示如何使用 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
from uni2ts.model.moirai_moe import MoiraiMoEForecast, MoiraiMoEModule

MODEL = "moirai2"  # 模型名称：可选 {'moirai', 'moirai-moe', 'moirai2'}
SIZE = "small"  # 模型尺寸：可选 {'small', 'base', 'large'}
PDT = 20  # 预测长度：任意正整数
CTX = 200  # 上下文长度：任意正整数
PSZ = "auto"  # patch size：可选 {"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 表示不重叠的窗口
)

# 通过从 Hugging Face 模型库下载模型权重来准备预训练模型
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,
    )
elif MODEL == "moirai2":
    model = Moirai2Forecast(
        module=Moirai2Module.from_pretrained(
            f"Salesforce/moirai-2.0-R-small",
        ),
        prediction_length=100,
        context_length=1680,
        target_dim=1,
        feat_dynamic_real_dim=0,
        past_feat_dynamic_real_dim=0,
    )

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 笔记本示例

更多常见任务的示例，例如可视化预测结果、从 Pandas DataFrame 进行预测等，请参阅 example 文件夹。

💻 命令行界面

我们提供了若干脚本，作为命令行界面，可轻松运行微调、评估，甚至预训练任务。配置文件由 Hydra 框架管理。

微调

首先，让我们看看如何使用 Uni2TS 在自定义数据集上微调预训练模型。Uni2TS 使用 Hugging Face 数据集库 处理数据加载，我们需要先将您的数据集转换为 Uni2TS 格式。如果您的数据集是简单的 Pandas DataFrame，我们可以使用以下脚本轻松处理。 本示例将使用流行的 长序列预测基准 中的 ETTh1 数据集。对于更复杂的用例，请参阅 此笔记本，了解如何在 Uni2TS 中使用自定义数据集的深入示例。有关基于原始配置的正式 LSF 微调实验，请参阅 此文件夹，其中包含 Shell 脚本和更详细的示例。

1. 首先，将您希望保存处理后数据集的目录路径添加到 .env 文件中。

echo "CUSTOM_DATA_PATH=PATH_TO_SAVE" >> .env

2. 运行以下脚本以将数据集处理为所需格式。对于 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'

在某些情况下，我们可能希望使用从训练集计算出的均值和标准差对数据进行归一化。这可以通过设置 --normalize 参数来实现。

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

3. 最后，我们可以使用适当的训练和验证数据配置文件，简单地运行微调脚本。预测配置，如补丁大小、上下文长度和预测长度，需由用户指定。由于 dataset_type 为 wide，因此 data.mode 被设置为 S，表示单变量设置。

python -m cli.train \
  -cp conf/finetune \
  exp_name=example_lsf \
  run_name=example_run \
  model=moirai_1.0_R_small \
  model.patch_size=32 \
  model.context_length=1000 \
  model.prediction_length=96 \
  data=etth1 \
  data.patch_size=32 \
  data.context_length=1000 \
  data.prediction_length=96 \
  data.mode=S \
  val_data=etth1 \
  val_data.patch_size=32 \
  val_data.context_length=1000 \
  val_data.prediction_length=96 \
  val_data.mode=S

评估

评估脚本可用于计算 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=PATH_TO_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 PATH_TO_SAVE
echo "LOTSA_V1_PATH=PATH_TO_SAVE" >> .env

然后，我们可以简单地运行以下脚本以启动预训练任务。 有关如何进一步自定义设置，请参阅相关文件和配置文件。

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

👀 引用

如果你在研究或应用中使用了本仓库，请使用以下 BibTeX 格式进行引用：

@article{liu2024moiraimoe,
  title={Moirai-MoE：利用稀疏专家混合模型赋能时间序列基础模型},
  author={刘旭、刘俊成、吴杰拉尔德、阿克苏·塔哈、梁宇轩、齐默尔曼·罗杰、刘程浩、萨瓦雷斯·西尔维奥、熊才明、萨胡·多延},
  journal={arXiv 预印本 arXiv:2410.10469},
  year={2024}
}

@article{aksu2024gifteval,
  title={GIFT-Eval：通用时间序列预测模型评估基准},
  author={阿克苏·塔哈、吴杰拉尔德、刘俊成、刘旭、刘程浩、萨瓦雷斯·西尔维奥、熊才明、萨胡·多延},
  journal={arXiv 预印本 arXiv:2410.10393},
  year={2024}
}

@inproceedings{woo2024moirai,
  title={统一训练通用时间序列预测 Transformer 模型},
  author={吴杰拉尔德、刘程浩、库马尔·阿克沙特、熊才明、萨瓦雷斯·西尔维奥、萨胡·多延},
  booktitle={第四十一届国际机器学习大会},
  year={2024}
}

伦理考量

本项目仅用于支持学术论文的研究目的。我们的模型、数据集和代码并未专门针对所有下游用途进行设计或评估。我们强烈建议用户在部署此模型之前，评估并解决与准确性、安全性和公平性相关的潜在问题。我们鼓励用户考虑人工智能的常见局限性，遵守适用法律，并在选择应用场景时遵循最佳实践，尤其是在高风险场景下——这些场景中，错误或滥用可能会对人们的生命、权利或安全造成重大影响。如需更多关于应用场景的指导，请参阅我们的《用户协议》和《人工智能用户协议》。

Uni2TS 快速上手指南

Uni2TS 是一个基于 PyTorch 的库，旨在为通用时间序列 Transformer 提供统一的预训练、微调、推理和评估框架。它包含了著名的 Moirai 和 Moirai-MoE 模型系列。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux 或 macOS (Windows 用户建议使用 WSL2)
• Python: 3.8 或更高版本
• 依赖管理: 推荐安装 virtualenv 用于创建隔离环境
• 硬件: 推荐使用 NVIDIA GPU 以加速推理和训练（CPU 亦可运行推理）

安装步骤

1. 克隆仓库

首先从 GitHub 克隆项目代码：

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

2. 创建虚拟环境

创建并激活 Python 虚拟环境：

virtualenv venv
. venv/bin/activate

(Windows 用户使用 . venv\Scripts\activate)

3. 安装依赖

您可以选择从源码安装（推荐，包含最新功能）或通过 PyPI 安装。

方式 A：从源码安装（支持 Notebook 示例）

pip install -e '.[notebook]'

方式 B：通过 PyPI 安装

pip install uni2ts

4. 配置环境变量

在项目根目录下创建 .env 文件（后续数据处理和自定义路径配置需要）：

touch .env

基本使用

以下示例展示如何使用 Uni2TS 加载预训练的 Moirai-2.0 模型进行零样本（Zero-shot）时间序列预测。该示例使用 pandas 读取数据，并利用 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
# 注意：根据实际安装的版本，导入路径可能略有不同，此处以 README 为准
from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
from uni2ts.model.moirai_moe import MoiraiMoEForecast, MoiraiMoEModule
# 假设已包含 Moirai2 相关导入 (若未直接导出，请参考 example 文件夹)
try:
    from uni2ts.model.moirai2 import Moirai2Forecast, Moirai2Module
except ImportError:
    # 兼容旧版本或特定构建，实际使用时请确保模型类可用
    Moirai2Forecast = None 
    Moirai2Module = None

# --- 配置参数 ---
MODEL = "moirai2"  # 模型选择: 'moirai', 'moirai-moe', 'moirai2'
SIZE = "small"     # 模型大小: 'small', 'base', 'large'
PDT = 20           # 预测长度 (prediction length)
CTX = 200          # 上下文长度 (context length)
PSZ = "auto"       # 补丁大小 (patch size): "auto", 8, 16, 32, 64, 128
BSZ = 32           # 批次大小 (batch size)
TEST = 100         # 测试集长度

# 1. 读取数据 (示例使用远程 CSV)
url = (
    "https://gist.githubusercontent.com/rsnirwan/c8c8654a98350fadd229b00167174ec4"
    "/raw/a42101c7786d4bc7695228a0f2c8cea41340e18f/ts_wide.csv"
)
df = pd.read_csv(url, index_col=0, parse_dates=True)

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

# 3. 划分训练集/测试集 (最后 TEST 个时间步作为测试集)
train, test_template = split(ds, offset=-TEST)

# 4. 构建滚动窗口评估数据
test_data = test_template.generate_instances(
    prediction_length=PDT,
    windows=TEST // PDT,
    distance=PDT, 
)

# 5. 加载预训练模型
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,
    )
elif MODEL == "moirai2":
    # Moirai-2.0 示例配置
    model = Moirai2Forecast(
        module=Moirai2Module.from_pretrained("Salesforce/moirai-2.0-R-small"),
        prediction_length=100,
        context_length=1680,
        target_dim=1,
        feat_dynamic_real_dim=0,
        past_feat_dynamic_real_dim=0,
    )

# 6. 创建预测器并执行预测
predictor = model.create_predictor(batch_size=BSZ)
forecasts = predictor.predict(test_data.input)

# 7. 可视化结果
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 支持 wide (宽表), long (长表) 和 wide_multivariate 格式。对于自定义数据集，可使用 uni2ts.data.builder.simple 模块进行预处理。
• 微调 (Fine-tuning): 库提供了基于 Hydra 配置的命令行接口 (cli.train)，可轻松在自定义数据集上微调模型。
• 更多示例: 请查看项目中的 example 文件夹获取完整的 Jupyter Notebook 教程，包括数据可视化、Pandas DataFrame 直接预测等场景。