在 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.pydata.org/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 的採用。感謝您關注我們的部落格，希望您發現這些資訊對您的專案有益。