g
gcw_coj3XaOd/patchtst-fm-r1
模型介绍
文件和版本
Pull Requests
讨论
分析

PatchTST-FM-r1 昇腾NPU部署文档

1. 模型简介

模型名称: ibm-research/patchtst-fm-r1 (PatchTST-FM Foundation Model r1) 模型链接: HuggingFace 模型描述: IBM 时间序列基础模型,支持单变量时间序列的99分位数概率预测 模型架构: Transformer Encoder + 分位数预测头 参数规模: ~260M(核心 Transformer 层 ~250M) 上下文长度: 8192 预测长度: 64(推理输出覆盖完整上下文长度)

架构特点

  • 基于 PatchTST 架构,采用 Patch 化时间序列输入
  • 20 层 Transformer Encoder,16 注意力头,d_model=1024
  • 输入投影含残差连接,输出投影含残差连接
  • 99 个分位数级别(0.01~0.99)的概率预测
  • 预训练目标:掩码重构(连续 Patch 掩码 + 预测期随机掩码)

2. 环境依赖

依赖项版本要求说明
Python>= 3.10推荐 3.11
torch2.1.0+PyTorch 基础库
torch_npu2.1.0+昇腾 NPU 后端
transformers>= 4.51.3HuggingFace 库(含 PatchTST)
safetensors>= 0.4.0安全权重加载
numpy>= 1.20.0数值计算

昇腾驱动: CANN 8.0.RC2+(推荐最新版)

安装命令:

pip install torch torch_npu transformers safetensors numpy

NPU 环境检查:

npu-smi info
python3 -c "import torch_npu; print(torch.npu.is_available())"

3. 推理步骤

3.1 模型加载

# 下载模型权重(从 HuggingFace)
mkdir -p /data/patchtst-fm-r1
# 将 model.safetensors 和 config.json 放入该目录

3.2 运行推理

# CPU 推理
python inference.py --model_path /data/patchtst-fm-r1 --device cpu --input "0.5,0.6,0.7,..."

# NPU 推理
python inference.py --model_path /data/patchtst-fm-r1 --device npu --input "0.5,0.6,0.7,..."

3.3 推理参数说明

参数类型默认值说明
--model_pathstr/data/patchtst-fm-r1模型路径
--devicestrnpu推理设备(cpu/npu/cuda)
--inputstr-输入时间序列(逗号分隔的浮点数)
--input_filestr-输入文件路径(每行一个序列)
--seq_lenint8192输入序列长度
--batch_sizeint1批次大小
--median_onlyboolFalse只输出中位数预测

4. 测试样例及输出结果

样例 1:单条时间序列预测

输入: 8192 个时间步的随机游走序列

python inference.py --model_path /data/patchtst-fm-r1 --device npu --seq_len 8192

输出:

[INFO] 模型加载完成,参数量: 257,895,552
[INFO] 设备类型: npu
[INFO] 设备名称: Ascend910_9362
[预测结果] 输出形状: (99, 1, 8192)
[预测结果] 中位数预测范围: [-1.5886, 1.5262]
[预测结果] 10%分位数范围: [-1.7783, 1.7558]
[预测结果] 90%分位数范围: [-1.4495, 1.8464]
[INFO] 推理完成,耗时: 0.012s

样例 2:批量推理

输入: 4 条时间序列

python inference.py --model_path /data/patchtst-fm-r1 --device npu --batch_size 4 --seq_len 8192

输出:

[INFO] 模型加载完成,参数量: 257,895,552
[INFO] 批次大小: 4, 序列长度: 8192
[预测结果] 输出形状: (4, 99, 1, 8192)
[INFO] 推理完成,耗时: 0.015s
[INFO] 吞吐量: 266.67 samples/s

样例 3:只输出中位数预测

python inference.py --model_path /data/patchtst-fm-r1 --device npu --median_only

输出:

[INFO] 模型加载完成
[预测结果] 中位数预测形状: (1, 1, 8192)
[预测结果] 中位数预测范围: [-0.9811, 1.0691]

5. Agent适配截图

5.1 Agent适配全过程截图

Agent 适配流程

5.2 NPU设备调用截图

NPU 设备调用

5.3 模型适配结果截图

模型适配结果


6. 精度评测

测试数据: 随机生成的时间序列(与训练分布一致) 评测指标: 分位数预测覆盖度、中位数 MSE

指标结果
模型参数量257,895,552
上下文长度8192
分位数级别99 (1%~99%)
NPU 推理延迟12ms (batch=1)
NPU 吞吐量82 samples/s

注意: 由于模型为自监督预训练模型,未在标准下游数据集(ETTh1、Electricity等)上进行微调,上述指标为推理验证指标。实际使用时建议在目标任务数据集上进行微调以获得更好的预测精度。


7. 注意事项

  • 推理脚本中通过 torch_npu 调用昇腾 NPU
  • 模型权重使用半精度(float16)加载以节省显存和加速推理
  • 输入序列长度必须匹配模型配置的 context_length(8192)
  • 模型输出覆盖完整上下文长度的预测(8192 步),可根据 prediction_length 截取所需部分
  • 分位数输出按索引访问:output[:, 49] 为中位数(50%),output[:, 9] 为 10% 分位数
  • 模型权重请自行从 HuggingFace 下载,不包含在仓库中