時系列基盤amazon chronos-t5をサンプルデータでfine tuningをする

初めに

過去に 比較的小さい時系列基盤モデルを触りました。こちらはfine tuingはできなかったので、fine tuingができてより大きいモデルを触っていきます

ayousanz.hatenadiary.jp

環境

• L4 GPU
• ubuntu22.04
• jupter notebook

ライブラリのインストール

# library install
!pip install git+https://github.com/amazon-science/chronos-forecasting.git

データの取得と分析

以下でデータを取得・分析をしていきます

import pandas as pd

# データ読み込み
# https://github.com/zhouhaoyi/ETDataset/blob/main/ETT-small/ETTh1.csv
df = pd.read_csv("ETTh1.csv")
print(len(df))
df.head(2)

以下のようにデータの中が表示されます

データ形式の変換

推論に用いれる形に変換をします

import torch

context_length = 512
forecast_horizon = 96

# データセット分割
df_train = df.iloc[-(context_length+forecast_horizon):-forecast_horizon]
df_test = df.iloc[-forecast_horizon:]

# 形式の変更
train_tensor = torch.tensor(df_train[["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"]].values, dtype=torch.float)
train_tensor = train_tensor.t()
test_tensor = torch.tensor(df_test[["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"]].values, dtype=torch.float)
test_tensor = test_tensor.t()

推論

モデルのロード

以下でモデルをロードします

import pandas as pd
import torch
from chronos import ChronosPipeline

pipeline = ChronosPipeline.from_pretrained(
    "amazon/chronos-t5-large",
    device_map="cuda",  # use "cpu" for CPU inference and "mps" for Apple Silicon
    torch_dtype=torch.bfloat16,
)

推論実行及びグラフにプロット

forecast = pipeline.predict(train_tensor, forecast_horizon, limit_prediction_length=False)
forecast_median_tensor, _ = torch.median(forecast, dim=1)

import matplotlib.pyplot as plt

channel_idx = 6
time_index = 0

history = train_tensor[channel_idx, :].detach().numpy()
true = test_tensor[channel_idx, :].detach().numpy()
pred = forecast_median_tensor[channel_idx, :].detach().numpy()

plt.figure(figsize=(12, 4))

# Plotting the first time series from history
plt.plot(range(len(history)), history, label='History (512 timesteps)', c='darkblue')

# Plotting ground truth and prediction
num_forecasts = len(true)

offset = len(history)
plt.plot(range(offset, offset + len(true)), true, label='Ground Truth (96 timesteps)', color='darkblue', linestyle='--', alpha=0.5)
plt.plot(range(offset, offset + len(pred)), pred, label='Forecast (96 timesteps)', color='red', linestyle='--')

plt.title(f"ETTh1 (Hourly) -- (idx={time_index}, channel={channel_idx})", fontsize=18)
plt.xlabel('Time', fontsize=14)
plt.ylabel('Value', fontsize=14)
plt.legend(fontsize=14)
plt.show()

以下のようにゼロショットで推論ができました。かなり一致しています

fine tuing

ライブラリのインストール

!pip install "chronos[training] @ git+https://github.com/amazon-science/chronos-forecasting.git"
!git clone https://github.com/amazon-science/chronos-forecasting.git

# Clone the ibm/tsfm
! git clone https://github.com/IBM/tsfm.git

# Change directory. Move inside the tsfm repo.
%cd tsfm

# Install the tsfm library
! pip install ".[notebooks]"

学習用のデータの変換

from pathlib import Path
from typing import List, Optional, Union

import numpy as np
from gluonts.dataset.arrow import ArrowWriter

from pathlib import Path
from typing import List, Optional, Union

import numpy as np
from gluonts.dataset.arrow import ArrowWriter


def convert_to_arrow(
    path: Union[str, Path],
    time_series: Union[List[np.ndarray], np.ndarray],
    start_times: Optional[Union[List[np.datetime64], np.ndarray]] = None,
    compression: str = "lz4",
):
    if start_times is None:
        # Set an arbitrary start time
        start_times = [np.datetime64("2000-01-01 00:00", "s")] * len(time_series)

    assert len(time_series) == len(start_times)

    dataset = [
        {"start": start, "target": ts} for ts, start in zip(time_series, start_times)
    ]
    ArrowWriter(compression=compression).write_to_file(
        dataset,
        path=path,
    )


# Convert to GluonTS arrow format
cols = ["HUFL", "HULL", "MUFL", "MULL", "LUFL", "LULL", "OT"]
convert_to_arrow(
    path = "./etth1-train.arrow", 
    time_series=[np.array(df_train[col]) for col in cols],
    start_times=[pd.to_datetime(df_train["date"]).values[0]] * len(cols),
)

以下のようなファイルが生成されます

参考サイト

hamaruki.com

note.com

fine tuing

以下で追加学習をします

import yaml

batch_size = 2  # バッチサイズ
num_steps = train_size/batch_size
print("steps:" + str(num_steps))

# Fine Tuningの設定
config_data = {
    'training_data_paths': ["./etth1-train.arrow"],  # 学習データファイルのパス
    'probability': [1.0],
    'output_dir': './output/',  # 学習結果の出力ディレクトリ
    'context_length': context_length,
    'prediction_length': forecast_horizon,
    'max_steps': num_steps,
    'per_device_train_batch_size': batch_size,
    'learning_rate': 0.001,
    'model_id': 'amazon/chronos-t5-large',
    # 'model_id': 'amazon/chronos-t5-base',
    'random_init': False,  # 事前学習済みモデルを使用
    'tf32': True,        # NVIDIA GPUの場合Trueにする
    'gradient_accumulation_steps':2,
}

# 設定ファイルをYAML形式で保存
config_file_path = './ft_config.yaml'
with open(config_file_path, 'w') as file:
    yaml.dump(config_data, file)

def fine_tune_chronos(train_file_path, config_file_path):
    """
    chronos-t5モデルをFine Tuningする関数

    Args:
        train_file_path (str): 学習用データファイルのパス
        config_file_path (str): Fine Tuning設定ファイルのパス
    """

    # Fine Tuningの実行
    !CUDA_VISIBLE_DEVICES=0 python chronos-forecasting/scripts/training/train.py --config {config_file_path}

# Fine Tuningの実行
fine_tune_chronos("./etth1-train.arrow", config_file_path)

ログは以下のようになりました

{'loss': 3.8517, 'grad_norm': 1.0407471656799316, 'learning_rate': 0.0009282433983926521, 'epoch': 0.07}
{'loss': 3.7261, 'grad_norm': 0.7981559038162231, 'learning_rate': 0.0008564867967853042, 'epoch': 0.14}
{'loss': 3.6565, 'grad_norm': 1.0834617614746094, 'learning_rate': 0.0007847301951779565, 'epoch': 0.22}
{'loss': 3.5609, 'grad_norm': 0.6380050182342529, 'learning_rate': 0.0007129735935706086, 'epoch': 0.29}
{'loss': 3.4786, 'grad_norm': 1.107723355293274, 'learning_rate': 0.0006412169919632607, 'epoch': 0.36}
{'loss': 3.3614, 'grad_norm': 0.8139169812202454, 'learning_rate': 0.0005694603903559128, 'epoch': 0.43}
{'loss': 3.2374, 'grad_norm': 0.9819308519363403, 'learning_rate': 0.0004977037887485649, 'epoch': 0.5}
{'loss': 3.1432, 'grad_norm': 1.3025404214859009, 'learning_rate': 0.000425947187141217, 'epoch': 0.57}
{'loss': 3.0098, 'grad_norm': 1.464908242225647, 'learning_rate': 0.0003541905855338691, 'epoch': 0.65}
{'loss': 2.8967, 'grad_norm': 1.4902337789535522, 'learning_rate': 0.00028243398392652127, 'epoch': 0.72}
{'loss': 2.7994, 'grad_norm': 1.2662851810455322, 'learning_rate': 0.00021067738231917335, 'epoch': 0.79}
{'loss': 2.7279, 'grad_norm': 1.741388201713562, 'learning_rate': 0.0001389207807118255, 'epoch': 0.86}
{'loss': 2.6553, 'grad_norm': 1.594014286994934, 'learning_rate': 6.716417910447761e-05, 'epoch': 0.93}
{'train_runtime': 7191.7126, 'train_samples_per_second': 3.876, 'train_steps_per_second': 0.969, 'train_loss': 3.195165220275949, 'epoch': 1.0}
100%|█████████████████████████████████████| 6968/6968 [1:59:51<00:00,  1.03s/it]

追加学習モデルを使った推論

from chronos import ChronosPipeline
import matplotlib.pyplot as plt

def predict_with_chronos(train_tensor, forecast_horizon, model_name="amazon/chronos-t5-large", device_map="cuda"):
    """
    chronos-t5モデルで予測を行う関数

    Args:
        train_tensor (torch.Tensor): 学習用データテンソル
        forecast_horizon (int): 予測する長さ
        model_name (str): モデル名 (デフォルト: "amazon/chronos-t5-large")
        device_map (str): デバイス ("cuda" or "cpu")

    Returns:
        forecast_median_tensor (torch.Tensor): 予測結果の中央値テンソル
    """

    pipeline = ChronosPipeline.from_pretrained(
        model_name,
        device_map=device_map,
        torch_dtype=torch.bfloat16,  # 計算精度を指定
    )

    # 予測の実行 (limit_prediction_length=Falseで予測長を制限しない)
    forecast = pipeline.predict(train_tensor, forecast_horizon, limit_prediction_length=False)
    forecast_median_tensor, _ = torch.median(forecast, dim=1)  # 予測結果の中央値を計算

    return forecast_median_tensor, forecast

def visualize_prediction(train_tensor, test_tensor, forecast_median_tensor, channel_idx=6):
    """
    予測結果を可視化する関数

    Args:
        train_tensor (torch.Tensor): 学習用データテンソル
        test_tensor (torch.Tensor): テスト用データテンソル
        forecast_median_tensor (torch.Tensor): 予測結果の中央値テンソル
        channel_idx (int): 可視化するチャンネルのインデックス (デフォルト: 6, OT)
    """

    history = train_tensor[channel_idx, :].detach().numpy()  # 学習用データ
    true = test_tensor[channel_idx, :].detach().numpy()      # 実測値
    pred = forecast_median_tensor[channel_idx, :].detach().numpy()  # 予測値

    plt.figure(figsize=(12, 4))
    plt.plot(range(len(history)), history, label='History (512 timesteps)', c='darkblue')
    plt.plot(range(len(history), len(history) + len(true)), true, label='Ground Truth (96 timesteps)', color='darkblue', linestyle='--', alpha=0.5)
    plt.plot(range(len(history), len(history) + len(pred)), pred, label='Forecast (96 timesteps)', color='red', linestyle='--')
    plt.title(f"ETTh1 (Hourly) - Channel {channel_idx}", fontsize=18)
    plt.xlabel('Time', fontsize=14)
    plt.ylabel('Value', fontsize=14)
    plt.legend(fontsize=14)
    plt.show()

# Fine Tuning後のモデルで予測
forecast_median_tensor_ft, forecast_ft = predict_with_chronos(train_tensor, forecast_horizon=forecast_horizon, model_name="./output/run-4/checkpoint-final/")

# 予測結果の可視化 (Fine Tuning後)
visualize_prediction(train_tensor, test_tensor, forecast_median_tensor_ft)