chronos-2 Ascend NPU 适配技能

# 安装依赖
pip install torch==2.9.0+cpu
pip install torch_npu==2.9.0.post1
pip install transformers>=4.37.0
pip install chronos>=0.3.0

# 设置镜像 (国内环境)
export HF_ENDPOINT=https://hf-mirror.com

# CANN 优化环境变量
export ASCEND_SLOG_PRINT_TO_STDOUT=0
export ASCEND_GLOBAL_LOG_LEVEL=3
export ACL_OP_COMPILER_CACHE_MODE=enable
export ACL_OP_DEBUG_LEVEL=0
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

import torch
import torch_npu
from chronos import ChronosPipeline

# NPU 初始化
device = torch.device("npu:0")
torch.npu.set_device(device)

# 加载模型 (自动映射到 NPU)
pipeline = ChronosPipeline.from_pretrained(
    "amazon/chronos-t5-small",
    device_map="npu:0",
    torch_dtype=torch.float32,
)

import numpy as np

# 准备时间序列数据 (batch=1, seq_len=512)
context = np.cumsum(np.random.randn(1, 512), axis=1)
context_tensor = torch.tensor(context, dtype=torch.float32)

# 预测未来 64 步
forecast = pipeline.predict(context_tensor, prediction_length=64)
# forecast shape: (batch, num_samples, prediction_length)
print(forecast.shape)  # e.g., (1, 20, 64)

# 增大 batch size 是提升吞吐量的最有效手段
context = np.cumsum(np.random.randn(32, 512), axis=1)
context_tensor = torch.tensor(context, dtype=torch.float32)

forecast = pipeline.predict(context_tensor, prediction_length=32)
# 吞吐量可达 38+ samples/s

# 适合离线批量预测场景
forecast = pipeline.predict(
    context_tensor,
    prediction_length=32,   # 缩短预测长度
    num_samples=20,         # 保持不确定性估计
)
# 推荐 batch_size=32
# 实测吞吐量: 38.06 samples/s

# 适合实时推理场景
forecast = pipeline.predict(
    context_tensor,
    prediction_length=16,   # 最短预测长度
    num_samples=5,          # 减少采样数
)
# 单条延迟: ~280ms

import numpy as np
from scipy import stats

def validate_npu_accuracy(npu_pipeline, cpu_pipeline, test_data, runs=30):
    npu_means = []
    cpu_means = []

    for _ in range(runs):
        npu_out = npu_pipeline.predict(test_data, 64)
        cpu_out = cpu_pipeline.predict(test_data, 64)
        npu_means.append(npu_out.cpu().mean().item())
        cpu_means.append(cpu_out.mean().item())

    t_stat, p_value = stats.ttest_ind(cpu_means, npu_means)
    return p_value > 0.05  # p > 0.05 表示分布一致

# 正确
ts_cpu = torch.tensor(data, dtype=torch.float32)
pipeline.predict(ts_cpu, 64)

# 错误
ts_npu = torch.tensor(data, dtype=torch.float32).to("npu:0")
pipeline.predict(ts_npu, 64)  # 会报错