← 返回案例展示

TAKT 案例展示:Rust 时间序列预测(ETT 上的 DLinear 和 PatchTST)

两个标准的长期预测模型 DLinear(LTSF-Linear)和 PatchTST,用官方脚本在 ETT 数据集上训练,从 PyTorch 翻译为纯 Rust(不使用 BLAS,不依赖 ML 运行时),并构建为两个版本:这份 Rust 代码的原始构建,以及同一代码的 TAKT 构建。两个构建的预测结果逐位一致。本包提供两个构建的可直接运行的程序、原始版本的源代码、根据官方数据重建输入窗口的脚本、基准测试脚本以及我们的测量数据。

程序计算的内容

ETT 数据集测试集中的每个窗口(官方划分边界,features M,StandardScaler 在训练集上拟合):输入为过去 336 步 x 7 个通道,输出为接下来的 96 步 x 7 个通道,float32。ETTh1/ETTh2:2,785 个窗口;ETTm1/ETTm2:11,425 个窗口。所选数据集对应检查点的权重编译在程序中(--data)。

模型 数据集 测试集 MSE / MAE(全部窗口,tools/score.py) 官方训练运行,测试集
DLinear ETTh1 0.384144 / 0.404713 0.384144 / 0.404713
DLinear ETTh2 0.290098 / 0.353328 0.290098 / 0.353328
DLinear ETTm1 0.301222 / 0.344619 0.301222 / 0.344619
DLinear ETTm2 0.171852 / 0.267122 0.171852 / 0.267122
PatchTST ETTh1 0.385126 / 0.405953(前 2,688 个:0.381618 / 0.405088) 0.381618 / 0.405088
PatchTST ETTh2 0.274616 / 0.337234(前 2,688 个:0.274118 / 0.336000) 0.274117 / 0.336000

官方 PatchTST 测试循环会丢弃最后一个不满 128 个窗口的批次,因此有“前 2,688 个”的数值。Rust 翻译版与 PyTorch 并非逐位一致(求和顺序不同;与 PyTorch eager 的最大绝对差为 2e-6 至 2e-5);TAKT 构建与 Rust 原始版本逐位一致。

测量(TAKT 测试台)

测量日期:2026-10-03。AMD Threadripper PRO 5975WX(Zen 3),Linux。预留一个 8 核 CCX,其 SMT 兄弟线程保持空闲;单线程运行绑定到一个核心,八线程运行绑定到该 CCX 的 8 个核心。每个数据集先进行一对预热运行,然后按交替顺序运行 9 轮,取中位数。时间:由程序自身计时的预测循环(不包括读取窗口和写入预测结果;这两部分在两个构建中开销相同)。时钟周期:同一循环的用户态时钟周期(perf stat,所有线程求和)。两个构建共用同一个驱动程序,它像长期运行的服务一样把已释放的堆内存保留在进程中(glibc mallopt),以免每个窗口的临时缓冲区引发缺页异常。所测文件正是 bin/ 中的文件,测量使用 tools/run_bench.sh。

模型 数据集 窗口数 线程数 原始版本(ms) TAKT(ms) 加速 原始版本(百万周期) TAKT(百万周期) 加速(按周期) 一致轮数
DLinear ETTh1 2,785 1 788.4 111.0 7.10× 3,518 486 7.23× 9/9
DLinear ETTh2 2,785 1 785.0 108.4 7.24× 3,512 476 7.38× 9/9
DLinear ETTm1 11,425 1 3,223.9 447.8 7.20× 14,421 1,978 7.29× 9/9
DLinear ETTm2 11,425 1 3,228.4 442.6 7.29× 14,405 1,984 7.26× 9/9
DLinear ETTh1 2,785 8 103.3 14.5 7.11× 3,515 486 7.23× 9/9
DLinear ETTh2 2,785 8 103.3 14.2 7.26× 3,514 477 7.37× 9/9
DLinear ETTm1 11,425 8 423.3 58.6 7.22× 14,426 1,988 7.26× 9/9
DLinear ETTm2 11,425 8 422.1 58.9 7.17× 14,407 1,998 7.21× 9/9
PatchTST ETTh1 2,785 1 13,440.8 5,733.4 2.34× 59,973 25,507 2.35× 9/9
PatchTST ETTh2 2,785 1 13,484.3 5,749.2 2.35× 60,068 25,645 2.34× 9/9
PatchTST ETTh1 2,785 8 1,777.5 765.7 2.32× 60,424 25,955 2.33× 9/9
PatchTST ETTh2 2,785 8 1,785.3 769.3 2.32× 60,677 26,070 2.33× 9/9

DLinear:无论单线程还是八线程,按时间加速 7.1–7.3×,按时钟周期加速 7.2–7.4×。PatchTST:按时间和按时钟周期均为 2.3×。八线程时,两个构建消耗的总时钟周期与单线程时大致相同。若计入整个进程(包括读取 26–107 MB 的窗口数据和写入预测结果),DLinear 单线程提速 5.7–6.0×,八线程提速 3.0–3.1×,PatchTST 提速 2.3×。详细信息、原始样本和整个进程的时间见 measurements.json。

运行

python3 tools/make_windows.py --out data --targets   # 下载官方 ETT CSV(经 sha256 校验),
                                                     # 写入 data/windows_*_test.f32(以及目标值)
mkdir -p out
for d in ETTh1 ETTh2 ETTm1 ETTm2; do
  bin/ltsf-dlinear-takt --data $d --windows data/windows_${d}_test.f32 --threads 1 --out out/dlinear_$d.f32
done
for d in ETTh1 ETTh2; do
  bin/ltsf-patchtst-takt --data $d --windows data/windows_${d}_test.f32 --threads 1 --out out/patchtst_$d.f32
done
sha256sum -c expected.sha256                 # 窗口、目标值和预测结果与我们的测量一致
python3 tools/score.py data/targets_ETTh1_test.f32 out/dlinear_ETTh1.f32

tools/run_bench.sh -m dlinear -n 9 -c 2          # 两个构建交替运行,逐字节检查,核心 2
tools/run_bench.sh -m dlinear -n 9 -t 8 -c 2-9   # 在核心 2-9 上运行八线程
tools/run_bench.sh -m patchtst -n 9 -c 2

--threads T 接受任意 T >= 1;预测结果与 T 无关。输入和输出文件均为原始的小端 float32(n x 336 x 7 和 n x 96 x 7);tools/make_windows.py 只需要 numpy。

Linux x86-64(glibc)。程序针对 x86-64-v3 构建(AVX2、FMA、BMI2:Intel Haswell 及更新、AMD Zen 及更新的处理器);不使用 FMA 指令。PatchTST 调用系统 C 库中的 erff:如果所用 glibc 的 erff 与 glibc 2.39(Ubuntu 24.04)不同,其预测结果可能与 expected.sha256 不符,但原始构建与 TAKT 构建之间仍逐字节一致。

等价性

重新构建原始版本

cd source
RUSTFLAGS="-C target-cpu=x86-64-v3 --remap-path-prefix=$HOME/.rustup=. \
  --remap-path-prefix=$HOME/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f=." \
  cargo +1.96.0 build --release --locked
strip --strip-all -o ltsf-dlinear-original target/release/ltsf-dlinear

使用 rustup 安装的 Rust 1.96.0(含 rust-src 组件),在我们的机器上可逐位复现 bin/ltsf-dlinear-original 和 bin/ltsf-patchtst-original;在其他环境下,嵌入的路径会不同,但预测结果不变。

早前与 PyTorch 的对比(无法用本包复现)

2026-10-02,我们在同一台机器上、在不同条件下测量了相同的模型:未预留测试台(39 个逻辑 CPU 与其他任务共享);两个 Rust 版本均以共享库形式针对本机 CPU 构建,并在进程内调用;PyTorch 2.14 运行在 CPU 上,分别使用 eager 模式和 torch.compile(Inductor),批大小为 64 或 512。对每种实现,线程数(1、8 或 32)和批大小均按验证集上的最低延迟选定,然后对测试集计时(5 次取中位数)。

模型 数据集 原始 Rust TAKT 构建 最快的 PyTorch(Inductor) PyTorch / TAKT
DLinear ETTh1 28.7 ms(32 线程) 5.1 ms(32) 18.5 ms(32,批大小 512) 3.63×
DLinear ETTh2 30.1 ms(32) 5.3 ms(32) 12.5 ms(32,批大小 512) 2.34×
DLinear ETTm1 115.4 ms(32) 19.7 ms(32) 65.1 ms(8,批大小 512) 3.31×
DLinear ETTm2 114.9 ms(32) 20.4 ms(32) 49.5 ms(32,批大小 512) 2.43×
PatchTST ETTh1 593 ms(32) 236 ms(32) 634 ms(32,批大小 64) 2.69×
PatchTST ETTh2 519 ms(32) 234 ms(32) 551 ms(32,批大小 64) 2.35×

单线程时排名不同:使用 Inductor 和批处理的 PyTorch 比 TAKT 构建更快(DLinear ETTh1:37.9 ms 对 111.4 ms;PatchTST ETTh1:2.03 s 对 5.78 s)。在那次运行中,在选定的 32 线程下,TAKT 构建比原始 Rust 快 5.6–5.9×(DLinear)和 2.2–2.5×(PatchTST),单线程下分别快 7.2–7.3× 和 2.3×。那次运行的全部数字见 measurements.json(earlier_measurement_2026_10_02)。

包内容

本包不含 ETT 数据(CC BY-ND 4.0,由 tools/make_windows.py 下载)。TAKT 构建的源代码不公开:我们交付的是构建产物。

原文: README.zh.md

Telegram