Chronos-2 Small (112M)一个时间序列预测的强大模型

功能对比

安装

pip install 'chronos-forecasting[extras]>=2.2' 'matplotlib'

快速开始

加载模型

import os

# 使用仅 1 个 GPU（如果可用）
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from chronos import BaseChronosPipeline, Chronos2Pipeline

# 加载 Chronos-2 管道
# 推荐使用 GPU 加速推理，CPU 也支持（使用 device_map="cpu"）
pipeline: Chronos2Pipeline = BaseChronosPipeline.from_pretrained(
    "amazon/chronos-2", 
    device_map="cuda"
)

单变量预测

# 加载数据（长格式 pandas DataFrame）
context_df = pd.read_csv("https://autogluon.s3.amazonaws.com/datasets/timeseries/m4_hourly/train.csv")
print("输入数据框形状:", context_df.shape)
display(context_df.head())

# 执行预测
pred_df = pipeline.predict_df(
    context_df, 
    prediction_length=24, 
    quantile_levels=[0.1, 0.5, 0.9]
)

print("输出数据框形状:", pred_df.shape)
display(pred_df.head())

predict_df 参数说明

带协变量的预测

# 能源价格预测配置
target = "target"  # 包含要预测的值的列名（能源价格）
prediction_length = 24  # 要预测的小时数
id_column = "id"  # 标识不同时间序列的列（国家/地区）
timestamp_column = "timestamp"  # 包含日期时间信息的列
timeseries_id = "DE"  # 要可视化的特定时间序列（德国）

# 加载历史能源价格和协变量的过去值
energy_context_df = pd.read_parquet(
    "https://autogluon.s3.amazonaws.com/datasets/timeseries/electricity_price/train.parquet"
)
energy_context_df[timestamp_column] = pd.to_datetime(energy_context_df[timestamp_column])

# 加载协变量的未来值
energy_test_df = pd.read_parquet(
    "https://autogluon.s3.amazonaws.com/datasets/timeseries/electricity_price/test.parquet"
)
energy_test_df[timestamp_column] = pd.to_datetime(energy_test_df[timestamp_column])
energy_future_df = energy_test_df.drop(columns=target)

# 使用协变量进行预测
pred_df = pipeline.predict_df(
    df=energy_context_df,
    future_df=energy_future_df,
    id_column=id_column,
    timestamp_column=timestamp_column,
    target=target,
    prediction_length=prediction_length,
    quantile_levels=[0.1, 0.5, 0.9],
)

可视化预测结果

def plot_forecast(
    context_df,
    test_df,
    pred_df,
    timeseries_id,
    target,
    timestamp_column="timestamp",
    id_column="item_id",
):
    fig, ax = plt.subplots(figsize=(12, 3))
    
    # 筛选特定时间序列的数据
    ts_context = context_df[context_df[id_column] == timeseries_id].set_index(timestamp_column)[target]
    ts_ground_truth = test_df[test_df[id_column] == timeseries_id].set_index(timestamp_column)[target]
    ts_pred = pred_df[pred_df[id_column] == timeseries_id].set_index(timestamp_column)
    
    # 绘制历史数据、真实值和预测值
    ts_context.plot(ax=ax, label=f"历史 {target}", color="xkcd:azure")
    ts_ground_truth.plot(ax=ax, label=f"未来 {target} (真实值)", color="xkcd:grass green")
    ts_pred["predictions"].plot(ax=ax, label="预测", color="xkcd:violet")
    
    # 绘制预测区间
    ax.fill_between(
        ts_pred.index,
        ts_pred["0.1"],
        ts_pred["0.9"],
        alpha=0.7,
        label="预测区间",
        color="xkcd:light lavender",
    )
    
    ax.axvline(x=ts_context.index[-1], color="black", linestyle="--", alpha=0.5)
    ax.legend(loc="upper left")
    ax.set_title(f"{target} 预测 - {timeseries_id}")
    fig.show()

支持的场景

硬件要求

参考