在 Hugging Face 中使用 PatchTSMixer - 入门指南

PatchTSMixer 是一种基于 MLP-Mixer 架构的轻量级时间序列建模方法。它由 IBM Research 的作者 Vijay Ekambaram、Arindam Jati、Nam Nguyen、Phanwadee Sinthong 和 Jayant Kalagnanam 在论文 TSMixer：用于多元时间序列预测的轻量级 MLP-Mixer 模型 中提出。

为了有效地分享思想并促进开源，IBM Research 与 HuggingFace 团队携手，在 Transformers 库中发布了该模型。

在 Hugging Face 的实现中，我们提供了 PatchTSMixer 的能力，可以轻松地在补丁、通道和隐藏特征之间进行轻量级混合，以实现有效的多元时间序列建模。它还支持各种注意力机制，从简单的门控注意到更复杂的自注意力块，这些都可以根据需要进行定制。该模型可以进行预训练，并随后用于各种下游任务，如预测、分类和回归。

PatchTSMixer 在预测方面以 8-60% 的显著优势优于最先进的 MLP 和 Transformer 模型。它还以显著减少的内存和运行时间（2-3 倍）优于最新的强大基准 Patch-Transformer 模型（优势为 1-2%）。更多详情，请参阅论文。

在这篇博客中，我们将演示如何开始使用 PatchTSMixer。我们首先将在 Electricity 数据集上展示 PatchTSMixer 的预测能力。然后，我们将通过使用在 Electricity 上训练的模型在 ETTH2 数据集上进行零样本预测，来展示 PatchTSMixer 的迁移学习能力。

PatchTSMixer 快速概览

如果您熟悉 PatchTSMixer，请跳过此部分！

PatchTSMixer 将给定的输入多元时间序列分割成一系列补丁或窗口。随后，它将序列传递给一个嵌入层，该层生成一个多维张量。

这个多维张量随后被传递到 PatchTSMixer 的主干网络，该主干网络由一系列 MLP Mixer 层组成。每个 MLP Mixer 层通过一系列置换和 MLP 操作学习补丁间、补丁内和通道间的相关性。

PatchTSMixer 还采用了残差连接和门控注意力机制来优先处理重要特征。

因此，一系列 MLP Mixer 层构成了以下的 PatchTSMixer 主干网络。

PatchTSMixer 采用模块化设计，无缝支持掩码时间序列预训练以及直接时间序列预测。

安装

此演示需要 Hugging Face 的 Transformers 库用于模型，以及 IBM 的 tsfm 包用于辅助数据预处理。两者都可以通过以下步骤安装。

1. 安装 IBM 时间序列基础模型库 tsfm。

pip install git+https://github.com/IBM/tsfm.git

2. 安装 Hugging Face Transformers

pip install transformers

3. 在 python 终端中使用以下命令进行测试。

from transformers import PatchTSMixerConfig
from tsfm_public.toolkit.dataset import ForecastDFDataset

第一部分：在 Electricity 数据集上进行预测

在这里，我们直接在 Electricity 数据集上训练一个 PatchTSMixer 模型，并评估其性能。

import os
import random

from transformers import (
    EarlyStoppingCallback,
    PatchTSMixerConfig,
    PatchTSMixerForPrediction,
    Trainer,
    TrainingArguments,
)
import numpy as np
import pandas as pd
import torch

from tsfm_public.toolkit.dataset import ForecastDFDataset
from tsfm_public.toolkit.time_series_preprocessor import TimeSeriesPreprocessor
from tsfm_public.toolkit.util import select_by_index

设置随机种子

from transformers import set_seed

set_seed(42)

加载和准备数据集

在下一个单元格中，请根据您的应用调整以下参数：

• dataset_path：本地 .csv 文件的路径，或目标数据的 csv 文件网址。数据使用 pandas 加载，因此任何 pd.read_csv 支持的格式都受支持：(https://pandas.ac.cn/pandas-docs/stable/reference/api/pandas.read_csv.html)。
• timestamp_column：包含时间戳信息的列名，如果没有此列，则使用 None。
• id_columns：指定不同时间序列 ID 的列名列表。如果不存在 ID 列，则使用 []。
• forecast_columns：要建模的列的列表。
• context_length：用作模型输入的历史数据量。将从输入数据框中提取长度等于 context_length 的输入时间序列数据窗口。在多时间序列数据集中，将创建上下文窗口，使其包含在单个时间序列（即单个 ID）内。
• forecast_horizon：要预测的未来时间戳数量。
• train_start_index, train_end_index：加载数据中用于划分训练数据的起始和结束索引。
• valid_start_index, valid_end_index：加载数据中用于划分验证数据的起始和结束索引。
• test_start_index, test_end_index：加载数据中用于划分测试数据的起始和结束索引。
• num_workers：PyTorch 数据加载器中的 CPU 工作进程数。
• batch_size：批次大小。数据首先加载到 Pandas 数据框中，并分为训练、验证和测试部分。然后，Pandas 数据框将转换为训练所需的相应 PyTorch 数据集。

# Download ECL data from https://github.com/zhouhaoyi/Informer2020
dataset_path = "~/Downloads/ECL.csv"
timestamp_column = "date"
id_columns = []

context_length = 512
forecast_horizon = 96
num_workers = 16  # Reduce this if you have low number of CPU cores
batch_size = 64  # Adjust according to GPU memory

data = pd.read_csv(
    dataset_path,
    parse_dates=[timestamp_column],
)
forecast_columns = list(data.columns[1:])

# get split
num_train = int(len(data) * 0.7)
num_test = int(len(data) * 0.2)
num_valid = len(data) - num_train - num_test
border1s = [
    0,
    num_train - context_length,
    len(data) - num_test - context_length,
]
border2s = [num_train, num_train + num_valid, len(data)]

train_start_index = border1s[0]  # None indicates beginning of dataset
train_end_index = border2s[0]

# we shift the start of the evaluation period back by context length so that
# the first evaluation timestamp is immediately following the training data
valid_start_index = border1s[1]
valid_end_index = border2s[1]

test_start_index = border1s[2]
test_end_index = border2s[2]

train_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=train_start_index,
    end_index=train_end_index,
)
valid_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=valid_start_index,
    end_index=valid_end_index,
)
test_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=test_start_index,
    end_index=test_end_index,
)

time_series_processor = TimeSeriesPreprocessor(
    context_length=context_length,
    timestamp_column=timestamp_column,
    id_columns=id_columns,
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    scaling=True,
)
time_series_processor.train(train_data)

train_dataset = ForecastDFDataset(
    time_series_processor.preprocess(train_data),
    id_columns=id_columns,
    timestamp_column="date",
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)
valid_dataset = ForecastDFDataset(
    time_series_processor.preprocess(valid_data),
    id_columns=id_columns,
    timestamp_column="date",
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)
test_dataset = ForecastDFDataset(
    time_series_processor.preprocess(test_data),
    id_columns=id_columns,
    timestamp_column="date",
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)

配置 PatchTSMixer 模型

接下来，我们使用一个配置实例化一个随机初始化的 PatchTSMixer 模型。以下设置控制与架构相关的不同超参数。

• num_input_channels：时间序列数据中的输入通道数（或维度）。这会自动设置为预测列的数量。
• context_length：如上所述，用作模型输入的历史数据量。
• prediction_length：这与上面描述的预测范围相同。
• patch_length：PatchTSMixer 模型的补丁长度。建议选择一个能被 context_length 整除的值。
• patch_stride：从上下文窗口中提取补丁时使用的步幅。
• d_model：模型的隐藏特征维度。
• num_layers：模型层数。
• dropout：编码器中所有全连接层的 Dropout 概率。
• head_dropout：模型头部使用的 Dropout 概率。
• mode：PatchTSMixer 的操作模式。"common_channel"/"mix_channel"。Common-channel 模式以通道独立的方式工作。对于预训练，请使用 "common_channel"。
• scaling：每个窗口的标准化缩放。推荐值："std"。

有关参数的完整详细信息，请参阅文档。

我们建议您只调整下一个单元格中的值。

patch_length = 8
config = PatchTSMixerConfig(
    context_length=context_length,
    prediction_length=forecast_horizon,
    patch_length=patch_length,
    num_input_channels=len(forecast_columns),
    patch_stride=patch_length,
    d_model=16,
    num_layers=8,
    expansion_factor=2,
    dropout=0.2,
    head_dropout=0.2,
    mode="common_channel",
    scaling="std",
)
model = PatchTSMixerForPrediction(config)

训练模型

接下来，我们可以利用 Hugging Face 的 Trainer 类，基于直接预测策略来训练模型。我们首先定义 TrainingArguments，其中列出了有关训练的各种超参数，如 epoch 数、学习率等。

training_args = TrainingArguments(
    output_dir="./checkpoint/patchtsmixer/electricity/pretrain/output/",
    overwrite_output_dir=True,
    learning_rate=0.001,
    num_train_epochs=100,  # For a quick test of this notebook, set it to 1
    do_eval=True,
    evaluation_strategy="epoch",
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    dataloader_num_workers=num_workers,
    report_to="tensorboard",
    save_strategy="epoch",
    logging_strategy="epoch",
    save_total_limit=3,
    logging_dir="./checkpoint/patchtsmixer/electricity/pretrain/logs/",  # Make sure to specify a logging directory
    load_best_model_at_end=True,  # Load the best model when training ends
    metric_for_best_model="eval_loss",  # Metric to monitor for early stopping
    greater_is_better=False,  # For loss
    label_names=["future_values"],
)

# Create the early stopping callback
early_stopping_callback = EarlyStoppingCallback(
    early_stopping_patience=10,  # Number of epochs with no improvement after which to stop
    early_stopping_threshold=0.0001,  # Minimum improvement required to consider as improvement
)

# define trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    callbacks=[early_stopping_callback],
)

# pretrain
trainer.train()

>>> | Epoch | Training Loss | Validation Loss |
    |-------|---------------|------------------|
    |   1   |    0.247100   |     0.141067     |
    |   2   |    0.168600   |     0.127757     |
    |   3   |    0.156500   |     0.122327     |
    ...

在测试集上评估模型

请注意，PatchTSMixer 的训练和评估损失是均方误差（MSE）损失。因此，在接下来的任何评估实验中，我们不会单独计算 MSE 指标。

results = trainer.evaluate(test_dataset)
print("Test result:")
print(results)

>>> Test result:
    {'eval_loss': 0.12884521484375, 'eval_runtime': 5.7532, 'eval_samples_per_second': 897.763, 'eval_steps_per_second': 3.65, 'epoch': 35.0}

我们得到的 MSE 分数为 0.128，这是 Electricity 数据上的 SOTA（最先进）结果。

保存模型

save_dir = "patchtsmixer/electricity/model/pretrain/"
os.makedirs(save_dir, exist_ok=True)
trainer.save_model(save_dir)

第二部分：从 Electricity 到 ETTh2 的迁移学习

在本节中，我们将展示 PatchTSMixer 模型的迁移学习能力。我们使用在 Electricity 数据集上预训练的模型，在 ETTh2 数据集上进行零样本预测。

迁移学习的意思是，我们首先在一个源数据集（我们上面在 Electricity 数据集上做的）上为预测任务预训练模型。然后，我们将在一个目标数据集上使用预训练模型进行零样本预测。零样本是指我们在目标领域测试性能而无需任何额外训练。我们希望模型从预训练中获得了足够的知识，可以迁移到不同的数据集上。

随后，我们将在目标数据的训练集上对预训练模型进行线性探测和（然后）微调，并将在目标数据的测试集上验证预测性能。在此示例中，源数据集是 Electricity 数据集，目标数据集是 ETTh2。

在 ETTh2 数据上进行迁移学习

所有评估均在 ETTh2 数据的测试部分进行： 步骤 1：直接评估在 electricity 上预训练的模型。这是零样本性能。
步骤 2：在进行线性探测后进行评估。
步骤 3：在进行全量微调后进行评估。

加载 ETTh2 数据集

下面，我们将 ETTh2 数据集加载为 Pandas 数据框。接下来，我们创建训练、验证和测试三个部分。然后，我们利用 TimeSeriesPreprocessor 类为模型准备每个部分。

dataset = "ETTh2"

dataset_path = f"https://raw.githubusercontent.com/zhouhaoyi/ETDataset/main/ETT-small/{dataset}.csv"
timestamp_column = "date"
id_columns = []
forecast_columns = ["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"]
train_start_index = None  # None indicates beginning of dataset
train_end_index = 12 * 30 * 24

# we shift the start of the evaluation period back by context length so that
# the first evaluation timestamp is immediately following the training data
valid_start_index = 12 * 30 * 24 - context_length
valid_end_index = 12 * 30 * 24 + 4 * 30 * 24

test_start_index = 12 * 30 * 24 + 4 * 30 * 24 - context_length
test_end_index = 12 * 30 * 24 + 8 * 30 * 24

data = pd.read_csv(
    dataset_path,
    parse_dates=[timestamp_column],
)

train_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=train_start_index,
    end_index=train_end_index,
)
valid_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=valid_start_index,
    end_index=valid_end_index,
)
test_data = select_by_index(
    data,
    id_columns=id_columns,
    start_index=test_start_index,
    end_index=test_end_index,
)

time_series_processor = TimeSeriesPreprocessor(
    context_length=context_length
    timestamp_column=timestamp_column,
    id_columns=id_columns,
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    scaling=True,
)
time_series_processor.train(train_data)

>>> TimeSeriesPreprocessor {
        "context_length": 512,
        "feature_extractor_type": "TimeSeriesPreprocessor",
        "id_columns": [],
    ...
    }

train_dataset = ForecastDFDataset(
    time_series_processor.preprocess(train_data),
    id_columns=id_columns,
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)
valid_dataset = ForecastDFDataset(
    time_series_processor.preprocess(valid_data),
    id_columns=id_columns,
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)
test_dataset = ForecastDFDataset(
    time_series_processor.preprocess(test_data),
    id_columns=id_columns,
    input_columns=forecast_columns,
    output_columns=forecast_columns,
    context_length=context_length,
    prediction_length=forecast_horizon,
)

在 ETTh2 上进行零样本预测

由于我们将测试开箱即用的预测性能，我们加载了上面预训练的模型。

from transformers import PatchTSMixerForPrediction

finetune_forecast_model = PatchTSMixerForPrediction.from_pretrained(
    "patchtsmixer/electricity/model/pretrain/"
)

finetune_forecast_args = TrainingArguments(
    output_dir="./checkpoint/patchtsmixer/transfer/finetune/output/",
    overwrite_output_dir=True,
    learning_rate=0.0001,
    num_train_epochs=100,
    do_eval=True,
    evaluation_strategy="epoch",
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    dataloader_num_workers=num_workers,
    report_to="tensorboard",
    save_strategy="epoch",
    logging_strategy="epoch",
    save_total_limit=3,
    logging_dir="./checkpoint/patchtsmixer/transfer/finetune/logs/",  # Make sure to specify a logging directory
    load_best_model_at_end=True,  # Load the best model when training ends
    metric_for_best_model="eval_loss",  # Metric to monitor for early stopping
    greater_is_better=False,  # For loss
)

# Create a new early stopping callback with faster convergence properties
early_stopping_callback = EarlyStoppingCallback(
    early_stopping_patience=5,  # Number of epochs with no improvement after which to stop
    early_stopping_threshold=0.001,  # Minimum improvement required to consider as improvement
)

finetune_forecast_trainer = Trainer(
    model=finetune_forecast_model,
    args=finetune_forecast_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    callbacks=[early_stopping_callback],
)

print("\n\nDoing zero-shot forecasting on target data")
result = finetune_forecast_trainer.evaluate(test_dataset)
print("Target data zero-shot forecasting result:")
print(result)

>>> Doing zero-shot forecasting on target data

    Target data zero-shot forecasting result:
    {'eval_loss': 0.3038313388824463, 'eval_runtime': 1.8364, 'eval_samples_per_second': 1516.562, 'eval_steps_per_second': 5.99}

可以看到，我们得到的零样本均方误差（MSE）为 0.3，接近最先进的结果。

接下来，让我们看看通过进行线性探测能达到什么效果，这涉及在冻结的预训练模型之上训练一个线性分类器。线性探测通常用于测试预训练模型的特征性能。

在 ETTh2 上进行线性探测

我们可以在目标数据的 train 部分上进行快速的线性探测，以查看是否有任何可能的 test 性能提升。

# Freeze the backbone of the model
for param in finetune_forecast_trainer.model.model.parameters():
    param.requires_grad = False

print("\n\nLinear probing on the target data")
finetune_forecast_trainer.train()
print("Evaluating")
result = finetune_forecast_trainer.evaluate(test_dataset)
print("Target data head/linear probing result:")
print(result)
    
>>> Linear probing on the target data


    | Epoch | Training Loss | Validation Loss |
    |-------|---------------|------------------|
    |   1   |    0.447000   |     0.216436     |
    |   2   |    0.438600   |     0.215667     |
    |   3   |    0.429400   |     0.215104     |
    ...

    Evaluating

    Target data head/linear probing result:
    {'eval_loss': 0.27119266986846924, 'eval_runtime': 1.7621, 'eval_samples_per_second': 1580.478, 'eval_steps_per_second': 6.242, 'epoch': 13.0}

可以看到，通过在冻结的主干网络之上训练一个简单的线性层，MSE 从 0.3 降至 0.271，达到了最先进的结果。

save_dir = f"patchtsmixer/electricity/model/transfer/{dataset}/model/linear_probe/"
os.makedirs(save_dir, exist_ok=True)
finetune_forecast_trainer.save_model(save_dir)

save_dir = f"patchtsmixer/electricity/model/transfer/{dataset}/preprocessor/"
os.makedirs(save_dir, exist_ok=True)
time_series_processor.save_pretrained(save_dir)

>>> ['patchtsmixer/electricity/model/transfer/ETTh2/preprocessor/preprocessor_config.json']

最后，让我们看看通过在目标数据集上对模型进行全量微调是否能获得更多改进。

在 ETTh2 上进行全量微调

我们可以在目标数据的 train 部分上进行全模型微调（而不是像上面那样只探测最后一个线性层），以查看可能的 test 性能提升。代码看起来与上面的线性探测任务类似，只是我们没有冻结任何参数。

# Reload the model
finetune_forecast_model = PatchTSMixerForPrediction.from_pretrained(
    "patchtsmixer/electricity/model/pretrain/"
)
finetune_forecast_trainer = Trainer(
    model=finetune_forecast_model,
    args=finetune_forecast_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    callbacks=[early_stopping_callback],
)
print("\n\nFinetuning on the target data")
finetune_forecast_trainer.train()
print("Evaluating")
result = finetune_forecast_trainer.evaluate(test_dataset)
print("Target data full finetune result:")
print(result)

>>> Finetuning on the target data

    | Epoch | Training Loss | Validation Loss |
    |-------|---------------|-----------------|
    |   1   |    0.432900   |     0.215200    |
    |   2   |    0.416700   |     0.210919    |
    |   3   |    0.401400   |     0.209932    |
    ...

    Evaluating

    Target data full finetune result:
    {'eval_loss': 0.2734043300151825, 'eval_runtime': 1.5853, 'eval_samples_per_second': 1756.725, 'eval_steps_per_second': 6.939, 'epoch': 9.0}

在这种情况下，通过全量微调并没有太多改进。不过我们还是保存模型吧。

save_dir = f"patchtsmixer/electricity/model/transfer/{dataset}/model/fine_tuning/"
os.makedirs(save_dir, exist_ok=True)
finetune_forecast_trainer.save_model(save_dir)

总结

在这篇博客中，我们提供了一个关于如何利用 PatchTSMixer 进行预测和迁移学习任务的逐步指南。我们旨在促进 PatchTSMixer HF 模型与您的预测用例的无缝集成。我们相信这些内容可以作为有用的资源，加快您对 PatchTSMixer 的采用。感谢您关注我们的博客，希望您发现这些信息对您的项目有益。