`0-finetune-chronos-etth.py` 代码说明

这份脚本在做什么

这份脚本做的事情其实可以用一句话概括：

用本地 chronos-t5-tiny 和本地 ETTh1.csv，跑一个最小可运行的时间序列微调闭环。

它不是完整研究级训练框架，而是一个课堂版 demo。
它的目标是让编程基础还不强的同学，也能看清楚下面这条路径：

先准备一段历史序列和对应未来序列
把它们切成很多小窗口
用 Chronos 自带 tokenizer 把数值变成 token
用 T5ForConditionalGeneration 做 teacher-forcing 训练
保存微调后的模型
比较微调前后的预测结果

如果你把它和前面：

5/2-prog/0-chronos-etth-demo.py

对照着看，会更容易理解：

2-prog：只做 zero-shot 预测
4-prog：开始做 最小微调

也就是说，这份脚本是在上一份 demo 的基础上往前走了一步。

先看整体结构

这份脚本可以分成 10 个部分：

导入依赖
定义一个数据样本结构 WindowExample
定义滑动窗口数据集 SlidingWindowDataset
读取 Chronos 配置，构造 tokenizer
定义 collate_fn
选择设备
定义一个“用 pipeline 做预测”的辅助函数
主函数里解析参数
主函数里做训练
主函数里做微调前后预测对比并保存结果

第 1 部分：导入依赖

import argparse
import json
import math
import os
from dataclasses import dataclass
from pathlib import Path

这些是 Python 标准库：

argparse 解析命令行参数
json 把训练结果写成 JSON 文件
math 这里主要用来处理 nan
os 用来清理代理环境变量
dataclass 用来定义简单数据结构
Path 用更清晰的方式处理文件路径

再看下面这一组：

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
from transformers import AutoConfig, AutoModelForSeq2SeqLM

from chronos import ChronosConfig, ChronosPipeline

它们的角色分别是：

matplotlib 画图
numpy 算平均值、分位数
pandas 读取 CSV
torch 深度学习训练的基础库
DataLoader, Dataset 组织训练数据
AdamW 优化器
AutoConfig 读取模型配置
AutoModelForSeq2SeqLM 加载 T5 结构模型
ChronosConfig 读取 Chronos 特有配置
ChronosPipeline 微调后做预测对比时使用的官方推理接口

第 2 部分：`WindowExample`

@dataclass
class WindowExample:
    context: torch.Tensor
    future: torch.Tensor

这个类非常简单，它只是把一个样本拆成两个部分：

context 历史窗口
future 对应的未来窗口

你可以把它理解成：

一个训练样本 = 过去一段 + 未来一段

为什么要这么写？

因为时间序列训练最自然的监督信号就是：

给模型过去，让它学会未来

第 3 部分：`SlidingWindowDataset`

这一段是整份脚本最重要的部分之一。

先看这个类在干什么

它把一整条长时间序列切成很多训练样本。

例如：

原始序列很长
我们每次拿 512 个点当 history
再拿后面 64 个点当 future

这样就能从一条长序列里切出很多 (context, future) 对。

`init` 里的参数是什么意思

def __init__(
    self,
    values: np.ndarray,
    context_length: int,
    prediction_length: int,
    max_windows: int,
    stride: int,
) -> None:

values 原始时间序列
context_length 每个样本里历史窗口的长度
prediction_length 每个样本里未来窗口的长度
max_windows 最多切多少个窗口，防止课堂 demo 太慢
stride 滑动步长。比如 8 表示每次往前挪 8 个点再切下一个窗口

先看核心变量

self.examples = []
total = context_length + prediction_length
upper = len(values) - total + 1

self.examples = [] 用来存所有样本
total 一个样本总共需要多少点
upper 最后一个可切窗口的起点上界

最关键的 for 循环

for start in range(0, max(upper, 0), stride):

这表示：

从第 0 个点开始切
每次往前移动 stride
直到不能再切出完整窗口为止

取出 context 和 future

context = values[start : start + context_length]
future = values[start + context_length : start + total]

这是最标准的时间序列监督学习切法：

前一段做输入
后一段做标签

这个判断在防什么

if len(future) < prediction_length:
    break

作用是：

如果最后剩下的数据不够长，就不要再切了

否则最后一个样本会不完整。

为什么要转成 `torch.tensor`

self.examples.append(
    WindowExample(
        context=torch.tensor(context, dtype=torch.float32),
        future=torch.tensor(future, dtype=torch.float32),
    )
)

因为后面训练要用 PyTorch，所以这里先统一转成：

torch.Tensor

并且用：

dtype=torch.float32

表示这是浮点数时间序列。

这一句在做什么

if len(self.examples) >= max_windows:
    break

这一步很重要，尤其在课堂 demo 里。

它的作用是：

不让数据集无限变大
让训练控制在一个很小的规模内

也就是说，这个脚本追求的是：

先跑通
先看懂

而不是先追求最优结果。

`len` 和 `getitem`

def __len__(self) -> int:
    return len(self.examples)

def __getitem__(self, idx: int) -> WindowExample:
    return self.examples[idx]

这是 PyTorch Dataset 的标准接口：

__len__ 告诉 DataLoader 数据集有多大
__getitem__ 告诉 DataLoader 给我第 idx 个样本

第 4 部分：`build_tokenizer`

def build_tokenizer(model_dir: Path):
    cfg = AutoConfig.from_pretrained(model_dir)
    chronos_cfg = ChronosConfig(**cfg.chronos_config)
    return chronos_cfg.create_tokenizer(), chronos_cfg

这段代码非常关键，因为它说明：

Chronos 并不是“普通 T5 直接拿来预测数值”

它还带了一层 Chronos 专用配置：

数值如何量化
vocab 多大
context / prediction 长度默认是多少

逐行看：

AutoConfig.from_pretrained(model_dir) 读取模型目录里的 config.json
cfg.chronos_config 取出其中 Chronos 特有的配置字段
ChronosConfig(**cfg.chronos_config) 把字典变成一个 Chronos 配置对象
chronos_cfg.create_tokenizer() 根据这份配置创建 tokenizer

这里最重要的理解是：

训练脚本不是自己发明一套 token 规则
而是复用模型自带的 Chronos tokenizer 规则

第 5 部分：`build_collate_fn`

这一段是第二个最关键的地方。

为什么需要 `collate_fn`

因为 DataLoader 每次拿到的是一批 WindowExample，
但模型真正需要的是：

input_ids
attention_mask
labels

collate_fn 的作用就是：

把“时间序列样本”
变成“模型可训练的 token 批次”

先把 batch 里的张量堆起来

contexts = torch.stack([item.context for item in batch], dim=0)
futures = torch.stack([item.future for item in batch], dim=0)

这里的 torch.stack 表示：

把很多 1D 序列拼成一个 2D 张量

结果就会像：

contexts.shape = [batch_size, context_length]
futures.shape = [batch_size, prediction_length]

这一句最重要

input_ids, attention_mask, scale = tokenizer.context_input_transform(contexts)

这一步在做：

把历史数值序列转成 Chronos 的 token 序列

返回三个东西：

input_ids token 序列
attention_mask 哪些位置是真实输入
scale 这个样本的缩放因子

为什么要 scale？

因为 Chronos 不是直接把原始数值硬离散化，而是：

先做缩放
再量化

这样不同量级的序列更容易共享一套 token 空间。

labels 怎么来

labels, labels_mask = tokenizer.label_input_transform(futures, scale)

这一句表示：

用同样的 scale，把未来窗口也转成 token labels

这一步非常像：

文本 seq2seq 任务里把 target sentence 变成 target token ids

为什么要把一部分 label 改成 `-100`

labels = labels.clone()
labels[~labels_mask] = -100

在 Hugging Face 的 seq2seq 训练里：

labels == -100 表示这个位置不参与 loss 计算

所以这里是在说：

padding 位置不要算 loss

最终返回的是什么

return {
    "input_ids": input_ids.long(),
    "attention_mask": attention_mask.long(),
    "labels": labels.long(),
}

这就是标准的 Hugging Face 训练输入格式。

第 6 部分：`pick_device`

def pick_device() -> torch.device:
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")

这一段很朴素：

有 GPU 就用 GPU
没有就退回 CPU

对学生来说，这也是一个很好的工程习惯：

代码不要只写死在 GPU 上

第 7 部分：`forecast_with_pipeline`

这段函数的作用是：

用官方 ChronosPipeline
对同一段 history 做预测

为什么要单独写这个函数？

因为我们后面要做两次：

用微调前的模型预测
用微调后的模型预测

这样就能直接比较：

fine-tune 前后到底变了什么

逐行看

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

这里表示：

GPU 上更适合用 bfloat16
CPU 上就老老实实用 float32

再看：

pipeline = ChronosPipeline.from_pretrained(
    str(model_dir),
    device_map=device,
    dtype=dtype,
)

这表示：

从某个模型目录加载 Chronos pipeline

这里的 model_dir 可能是：

这正是这份脚本最漂亮的地方：

真正做预测的是这句

forecast = pipeline.predict(context, prediction_length, num_samples=num_samples)

这里：

context 是历史窗口
prediction_length 未来要预测多长
num_samples 采样多少条未来轨迹

返回结果不是一条序列，而是：

多个样本预测

所以后面我们才会再去算：

0.1 分位数
0.5 分位数
0.9 分位数

第 8 部分：`quantiles`

def quantiles(samples: np.ndarray):
    return np.quantile(samples, [0.1, 0.5, 0.9], axis=0)

这一段非常短，但很重要。

它做的是：

从多条样本预测里提取几个代表性的统计量

这里：

0.1 可以理解成偏低边界
0.5 是中位数预测
0.9 是偏高边界

这样就能画出：

预测区间

第 9 部分：`main()` 里的参数

这一部分先用 argparse 定义命令行参数。

基础资源参数

--base-model
--csv
--target

--base-model 基础模型路径，默认是 ../2-prog/models/chronos-t5-tiny
--csv 数据文件路径，默认是 ../2-prog/data/ETTh1.csv
--target 要预测的列，默认是 OT

窗口参数

--context-length
--prediction-length
--stride

--context-length 每个样本里历史窗口长度
--prediction-length 每个样本里未来窗口长度
--stride 相邻窗口起点之间相隔多少步

数据量控制参数

--max-train-windows
--max-val-windows

--max-train-windows 训练最多切多少个窗口
--max-val-windows 验证最多切多少个窗口

这两个参数最主要是为了：

控制课堂 demo 运行时间

训练参数

--epochs
--batch-size
--learning-rate

--epochs 训练多少轮
--batch-size 每次喂给模型多少个样本
--learning-rate 每次参数更新的步子有多大

推理与输出参数

--num-samples
--output-dir
--plot-path
--metrics-path

--num-samples 预测时采样几条未来轨迹
--output-dir 微调后模型保存到哪里
--plot-path 预测对比图保存到哪里
--metrics-path 指标 JSON 保存到哪里

第 10 部分：清理代理环境变量

for key in [...]:
    os.environ.pop(key, None)

这一段和你前面本地 Ollama demo 里看到的很像。

作用是：

避免环境里的代理设置干扰本地文件 / 本地加载流程

课堂上最简单的理解就是：

先把环境变量里的“网络弯路”拿掉

第 11 部分：把路径准备好

base_model = Path(args.base_model).resolve()
csv_path = Path(args.csv).resolve()
output_dir = Path(args.output_dir)
plot_path = Path(args.plot_path)
metrics_path = Path(args.metrics_path)

这一步就是：

把命令行参数变成路径对象

后面这几句：

output_dir.mkdir(parents=True, exist_ok=True)
plot_path.parent.mkdir(parents=True, exist_ok=True)
metrics_path.parent.mkdir(parents=True, exist_ok=True)

作用是：

这是一种很实用的工程习惯：

第 12 部分：基础检查

if not base_model.exists():
    raise FileNotFoundError(...)
if not csv_path.exists():
    raise FileNotFoundError(...)

这是在提前防止最常见错误：

模型目录没放好
数据文件没放好

比起让程序后面莫名报错，这种“提前检查并报清楚”更适合教学。

第 13 部分：构造 tokenizer 并检查预测长度

tokenizer, chronos_cfg = build_tokenizer(base_model)
if args.prediction_length != chronos_cfg.prediction_length:
    raise ValueError(...)

这个检查非常关键。

为什么？

因为当前 chronos-t5-tiny 的配置本身规定了：

prediction_length = 64

如果你硬把训练标签长度改成别的值，tokenizer 这边就不一致了。

所以这里是在提醒学生：

预训练模型不是任意参数都能随便改
模型配置和训练数据接口必须对齐

第 14 部分：读数据并选目标列

df = pd.read_csv(csv_path)
if args.target not in df.columns:
    raise ValueError(...)
values = df[args.target].astype(float).to_numpy()

逐行解释：

pd.read_csv(csv_path) 读入 ETTh1 数据
检查目标列是否存在避免学生把列名写错
astype(float) 强制转成浮点数
to_numpy() 从 DataFrame 列变成 NumPy 数组

为什么最后要变成 NumPy？

因为后面自己切窗口时，NumPy 一维数组最直接。

第 15 部分：划分 train / val / test

split_train = int(len(values) * 0.8)
split_val = int(len(values) * 0.9)

这里使用的是最简单的时间顺序切分：

前 80% 做训练
接下来 10% 做验证
最后 10% 做测试

这比随机打乱更符合时间序列场景，因为：

时间序列不能像普通分类那样随便打乱时间顺序

为什么 `val_values` 和 `test_values` 前面多切了一段

val_values = values[split_train - args.context_length - args.prediction_length : split_val]
test_values = values[split_val - args.context_length - args.prediction_length :]

这是为了保证：

在验证段和测试段里，也能切出完整窗口

也就是说，你不能只给验证集剩一小段未来，还得留足够长的历史上下文。

第 16 部分：构造 Dataset 和 DataLoader

train_ds = SlidingWindowDataset(...)
val_ds = SlidingWindowDataset(...)
collate_fn = build_collate_fn(tokenizer)
train_loader = DataLoader(...)
val_loader = DataLoader(...)

这一段的逻辑是：

先把原始长序列切成很多 (context, future) 样本
再把这些样本批量喂给模型

这里：

shuffle=True 训练时打乱窗口顺序
shuffle=False 验证时不打乱

这是 PyTorch 里非常常见的模式。

第 17 部分：加载模型

device = pick_device()
model = AutoModelForSeq2SeqLM.from_pretrained(str(base_model))
model.to(device)
model.train()

逐行看：

pick_device() 选 CPU 还是 GPU
AutoModelForSeq2SeqLM.from_pretrained(...) 加载底层 T5 模型
model.to(device) 把模型搬到对应设备
model.train() 切换到训练模式

这里最值得学生记住的是：

虽然外面名字叫 Chronos
但它底层其实是一个 seq2seq 语言模型

第 18 部分：优化器和损失记录

optimizer = AdamW(model.parameters(), lr=args.learning_rate)
train_losses = []
val_losses = []

AdamW 常见深度学习优化器
model.parameters() 告诉优化器：你要更新哪些参数
lr=args.learning_rate 学习率
train_losses / val_losses 记录每轮训练和验证损失

第 19 部分：训练循环

这一段是整份脚本最像“标准训练代码”的地方。

外层 epoch 循环

for epoch in range(args.epochs):

意思是：

训练若干轮

内层 batch 循环

for batch in train_loader:
    batch = {k: v.to(device) for k, v in batch.items()}
    outputs = model(**batch)
    loss = outputs.loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    epoch_losses.append(loss.item())

逐行解释：

for batch in train_loader 一批一批拿训练样本
batch = {k: v.to(device) ...} 把这一批张量搬到 CPU 或 GPU
outputs = model(**batch) 调用模型前向传播
loss = outputs.loss 直接取 Hugging Face 模型算好的损失
loss.backward() 反向传播，计算梯度
optimizer.step() 按梯度更新参数
optimizer.zero_grad() 清空旧梯度，避免累积
epoch_losses.append(loss.item()) 把这一批的损失存下来

这就是最标准的 PyTorch 训练闭环。

为什么还要验证

训练完一轮后，这段代码会跑验证集：

model.eval()
with torch.no_grad():
    ...
model.train()

这里：

model.eval() 切到评估模式
torch.no_grad() 不算梯度，省显存省时间
model.train() 验证完再切回训练模式

这一步的意义是：

看模型是不是只在训练集上越来越会背
还是在验证集上也有改善

第 20 部分：保存模型

model.save_pretrained(output_dir)
AutoConfig.from_pretrained(base_model).save_pretrained(output_dir)

第一句表示：

把训练后的参数存下来

第二句表示：

把配置文件也一起存下来

为什么配置也要存？

因为后面要重新用 ChronosPipeline.from_pretrained(output_dir) 加载这个模型。
如果只有权重、没有配置，很多框架就不知道该怎么还原它。

第 21 部分：取测试窗口

history = test_values[: args.context_length]
future = test_values[args.context_length : args.context_length + args.prediction_length]

这表示：

从测试段里取一段 history
再取后面紧接着的一段 future

它和训练时切窗口的思想完全一样，只不过这里我们只拿一个样本来做展示。

第 22 部分：比较微调前后预测

before = forecast_with_pipeline(base_model, ...)
after = forecast_with_pipeline(output_dir, ...)

这两句是整份脚本的教学高潮。

它们分别表示：

微调前模型怎么预测
微调后模型怎么预测

这一步让学生能非常直观地看到：

fine-tune 到底有没有把模型推向更适合当前数据的方向

第 23 部分：算 MAE

mae_before = float(np.mean(np.abs(med_b - future)))
mae_after = float(np.mean(np.abs(med_a - future)))

这里的 MAE 就是：

平均绝对误差

为什么拿中位数预测 med_b / med_a 去和真实值比？

因为：

预测输出是一个分布
中位数是最容易解释的代表性预测

第 24 部分：画图

plt.plot(x_hist, history, ...)
plt.plot(x_pred, future, ...)
plt.plot(x_pred, med_b, ...)
plt.fill_between(...)
plt.plot(x_pred, med_a, ...)
plt.fill_between(...)

这张图会同时画出：

历史
真实未来
微调前预测
微调后预测

所以它非常适合课堂上讲：

fine-tune 前后到底哪里变了

如果学生不太会看数字指标，这张图往往比 MAE 更直观。

第 25 部分：保存 metrics

metrics = {
    ...
}
metrics_path.write_text(json.dumps(metrics, indent=2))

这一段做的是：

把训练和评估的关键结果保存成 JSON 文件

为什么这一步重要？

因为真实实验里，只看终端输出是不够的。
你需要一个能复查的结果文件。

第 26 部分：最后打印结果

print(f"Saved fine-tuned model to: {output_dir}")
print(f"Saved comparison plot to: {plot_path}")
print(f"Saved metrics to: {metrics_path}")
print(f"MAE before fine-tune: {mae_before:.3f}")
print(f"MAE after fine-tune:  {mae_after:.3f}")

这一步很适合教学，因为它会把学生最关心的 5 件事一次性说清楚：

模型存到哪里了
图存到哪里了
指标存到哪里了
微调前误差是多少
微调后误差是多少

这份脚本最想让你理解什么

如果把整份脚本压缩成 4 句话，它真正想让你理解的是：

时间序列 foundation model 也可以像文本模型一样继续微调。
Chronos 的关键不是直接喂原始数值，而是先做 Chronos tokenizer 的数值 token 化。
最小微调闭环并不神秘，本质上还是“数据集 -> DataLoader -> 模型 -> loss -> optimizer”。
fine-tune 最后是否有价值，不看感觉，要看“微调前 vs 微调后”的结果对比。

如果你是第一次读这份代码，最该先盯住哪几段

如果你觉得整份脚本还是长，最推荐先盯住下面 5 个位置：

SlidingWindowDataset 看清楚时间序列样本是怎样切出来的
build_collate_fn 看清楚数值是怎样变成 token 的
AutoModelForSeq2SeqLM.from_pretrained(...) 看清楚 Chronos 底层其实是什么模型
训练循环看清楚 PyTorch 的最基本闭环
before / after 看清楚为什么最终一定要做微调前后对比

0-finetune-chronos-etth.py 代码说明

这份脚本在做什么

先看整体结构

第 1 部分：导入依赖

第 2 部分：WindowExample

第 3 部分：SlidingWindowDataset

先看这个类在干什么

__init__ 里的参数是什么意思

先看核心变量

最关键的 for 循环

取出 context 和 future

这个判断在防什么

为什么要转成 torch.tensor

这一句在做什么

__len__ 和 __getitem__

第 4 部分：build_tokenizer

第 5 部分：build_collate_fn

为什么需要 collate_fn

先把 batch 里的张量堆起来

这一句最重要

labels 怎么来

为什么要把一部分 label 改成 -100

最终返回的是什么

第 6 部分：pick_device

第 7 部分：forecast_with_pipeline

逐行看

真正做预测的是这句

第 8 部分：quantiles

第 9 部分：main() 里的参数

基础资源参数

窗口参数

数据量控制参数

训练参数

推理与输出参数

第 10 部分：清理代理环境变量

第 11 部分：把路径准备好

第 12 部分：基础检查

第 13 部分：构造 tokenizer 并检查预测长度

第 14 部分：读数据并选目标列

第 15 部分：划分 train / val / test

为什么 val_values 和 test_values 前面多切了一段

第 16 部分：构造 Dataset 和 DataLoader

第 17 部分：加载模型

第 18 部分：优化器和损失记录

第 19 部分：训练循环

外层 epoch 循环

内层 batch 循环

为什么还要验证

第 20 部分：保存模型

第 21 部分：取测试窗口

第 22 部分：比较微调前后预测

第 23 部分：算 MAE

第 24 部分：画图

第 25 部分：保存 metrics

第 26 部分：最后打印结果

这份脚本最想让你理解什么

如果你是第一次读这份代码，最该先盯住哪几段

`0-finetune-chronos-etth.py` 代码说明

第 2 部分：`WindowExample`

第 3 部分：`SlidingWindowDataset`

`init` 里的参数是什么意思

为什么要转成 `torch.tensor`

`len` 和 `getitem`

第 4 部分：`build_tokenizer`

第 5 部分：`build_collate_fn`

为什么需要 `collate_fn`

为什么要把一部分 label 改成 `-100`

第 6 部分：`pick_device`

第 7 部分：`forecast_with_pipeline`

第 8 部分：`quantiles`

第 9 部分：`main()` 里的参数

为什么 `val_values` 和 `test_values` 前面多切了一段