Time-MoE 开发者指南
1. 项目介绍
Time-MoE 是一个基于 Mixture of Experts(MoE)架构的时间序列基础模型,旨在处理大规模时间序列数据的预测任务。该项目是第一个将时间序列基础模型参数扩展到 24 亿规模的尝试,从零开始训练。Time-MoE 通过其独特的 MoE 架构,能够实现高效的计算和长序列的处理。
2. 项目快速启动
在开始之前,请确保安装了 Python 3.10 或更高版本。
安装依赖
pip install -r requirements.txt
Time-MoE 需要 transformers==4.40.1
,也可以选择安装 flash-attn
以提高训练和推理的速度:
pip install flash-attn==2.6.3
制作预测
以下是一个简单的预测示例:
import torch
from transformers import AutoModelForCausalLM
context_length = 12
seqs = torch.randn(2, context_length) # [batch_size, context_length]
model = AutoModelForCausalLM.from_pretrained('Maple728/TimeMoE-50M', device_map="cpu", trust_remote_code=True)
mean, std = seqs.mean(dim=-1, keepdim=True), seqs.std(dim=-1, keepdim=True)
normed_seqs = (seqs - mean) / std
prediction_length = 6
output = model.generate(normed_seqs, max_new_tokens=prediction_length)
normed_predictions = output[:, -prediction_length:]
predictions = normed_predictions * std + mean
如果序列已经归一化,可以跳过均值和标准差的计算。
3. 应用案例和最佳实践
评估模型
准备基准数据集,可以从 Google Drive 获取预处理好的数据集,并将其放在 ./dataset
目录下。以下是一个评估 ETTh1 数据集的示例:
python run_eval.py -d dataset/ETT-small/ETTh1.csv -p 96
微调 Time-MoE
要开始微调 Time-MoE,您的数据集需要转换为 jsonl
格式。以下是数据格式的一个例子:
{
"sequence": [1.0, 2.0, 3.0, ...]
}
{
"sequence": [11.0, 22.0, 33.0, ...]
}
可以使用 jsonl
、json
或 pickle
格式保存转换后的数据。如果使用 Time-300B 数据集,可以直接进行,无需额外预处理。
使用 CPU 进行训练的命令如下:
python main.py -d <data_path>
对于单节点单 GPU 或多 GPU 训练,使用以下命令:
python torch_dist_run.py main.py -d <data_path>
多节点多 GPU 训练需要设置环境变量以支持节点间的通信:
export MASTER_ADDR=<master_addr>
export MASTER_PORT=<master_port>
export WORLD_SIZE=<world_size>
export RANK=<rank>
python torch_dist_run.py main.py -d <data_path>
要从头开始训练 Time-MoE,只需在命令中添加 --from_scratch
参数。
4. 典型生态项目
Time-MoE 作为时间序列预测的开源项目,其生态系统包括但不限于以下项目:
- TimeMix:一个用于时间序列数据混洗的工具。
- 其他时间序列处理和预测框架,如 Prophet、LSTM 等。
通过上述介绍,开发者可以快速上手 Time-MoE,并在实际项目中应用时间序列预测的最佳实践。