PythonAI算法
Transformer 时间融合模型在股票价格预测中的应用与代码实现
本文介绍基于时间融合 Transformer（TFT）模型的股票价格预测方法。涵盖高频数据清洗、特征工程构建、TFT 架构原理及 PyTorch 实现细节。通过置信区间评估风险，结合量化指标优化交易策略，提供完整的机器学习建模流程与代码示例。重点讲解变量选择网络、时序注意力机制及训练评估流程，旨在为量化交易提供可解释的预测工具。
本文介绍基于时间融合 Transformer（TFT）模型的股票价格预测方法。涵盖高频数据清洗、特征工程构建、TFT 架构原理及 PyTorch 实现细节。通过置信区间评估风险，结合量化指标优化交易策略，提供完整的机器学习建模流程与代码示例。重点讲解变量选择网络、时序注意力机制及训练评估流程，旨在为量化交易提供可解释的预测工具。
股票市场是一个动态且充满不确定性的环境。成功的交易不仅依赖于对价格的预测，更在于对风险的有效评估与管理。传统的统计模型往往难以捕捉金融时间序列中的非线性特征和长期依赖关系。近年来，深度学习技术，特别是 Transformer 架构，在自然语言处理领域取得了巨大成功，并逐渐被引入到时间序列预测任务中。
本文详细介绍如何使用时间融合 Transformer（Time-Fusion Transformer, TFT）构建一个可解释的模型，用于股票价格的高频预测。TFT 模型不仅能输出预测值，还能提供置信区间，这对于量化交易中的风险管理至关重要。我们将涵盖从数据收集、清洗、特征工程到模型训练、评估及解释的完整流程。
免责声明：股票价格预测本质上具有高度不确定性。本教程仅供教育和研究目的使用，讨论的模型和方法不应被视为财务建议或实际交易的依据。在做出任何投资决策前，请务必进行独立研究并咨询专业金融顾问。
自动化交易算法能够帮助交易者消除情绪干扰，专注于技术指标和系统性决策。然而，追求完美的预测准确度在金融市场中几乎是不可能的。我们的目标不是实现 100% 的准确率，而是开发一个能够识别高概率变动趋势，并能量化每个预测相关风险的模型。
通过结合精确度预测与置信区间评估，该模型可以作为辅助交易决策的工具。TFT 模型的优势在于其可解释性，它允许我们观察哪些变量（如成交量、历史价格）对当前预测贡献最大，从而增强策略的可信度。
本项目使用了平均每日交易量超过 100 万股的股票数据，以确保流动性和数据的代表性。数据集涵盖了从 2024 年 1 月 1 日到 2024 年 7 月 11 日的 6 个月期间，包含超过 1500 只股票。对于每只股票，我们收集了标准的 OHLCV（开盘价、最高价、最低价、收盘价、成交量）数据，粒度为 1 分钟 K 线。
为了便于高效访问和分析，原始数据被保存为 .parquet 格式文件。
首先，导入必要的 Python 库，包括数据处理、可视化和建模工具：
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
加载数据集并将日期列设置为索引，以便进行基于时间的操作：
df = pd.read_parquet("data.parquet")
df['datetime'] = pd.to_datetime(df['datetime'])
df.set_index("datetime", inplace=True)
print(df.head())
检查数据的基本信息：
print(df.info())
输出显示数据框包含超过 6600 万行，占用约 4.5 GB 内存。为了降低计算复杂度并过滤掉流动性较差的股票，我们根据'平均名义成交量'进行筛选。名义成交量定义为日均成交量乘以平均价格。
def resample_to_daily(group):
resampled_group = group.resample('D').agg({
'open': 'first',
'high': 'max',
'low': 'min',
'close': 'last',
'volume': 'sum'
}).dropna()
return resampled_group
# 按股票代码分组并重采样为日线数据
df_daily = df.groupby('symbol').apply(resample_to_daily).reset_index()
# 计算统计指标以筛选高流动性股票
df_daily["stock_median_volume"] = df_daily.groupby("symbol")["volume"].transform("median").astype(int)
df_daily["stock_average_price"] = df_daily.groupby("symbol")["close"].transform("mean")
df_daily["stock_average_nominal_volume"] = df_daily["stock_median_volume"] * df_daily["stock_average_price"]
# 选取名义成交量最高的前 100 只股票
top_100_symbols = df_daily.groupby("symbol")["stock_average_nominal_volume"].median().nlargest(100).index
df_filtered = df[df['symbol'].isin(list(top_100_symbols))]
print(f"重采样后保留的 K 线数量：{df_filtered.shape[0]}")
print(f"涉及股票数量：{len(top_100_symbols)}")
在进行模型训练之前，深入理解数据分布和潜在模式至关重要。我们将进行以下特征工程步骤：
计算滚动窗口统计量，如移动平均线和波动率，这些是技术分析中的常用指标。
def add_features(df, window=20):
# 移动平均线
df['ma_' + str(window)] = df['close'].rolling(window=window).mean()
# 波动率 (标准差)
df['volatility'] = df['close'].rolling(window=window).std()
# 相对强弱指标 (RSI) 简化版
delta = df['close'].diff()
gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
rs = gain / loss
df['rsi'] = 100 - (100 / (1 + rs))
# 填充 NaN 值
df = df.fillna(method='bfill')
return df
df_processed = df_filtered.groupby('symbol').apply(add_features).reset_index(drop=True)
引入过去的时间步作为输入特征，帮助模型捕捉时间依赖性。
def create_lag_features(df, lags=[1, 5, 10]):
for lag in lags:
df[f'lag_close_{lag}'] = df.groupby('symbol')['close'].shift(lag)
return df
df_processed = df_processed.groupby('symbol').apply(create_lag_features).reset_index(drop=True)
定义预测目标。这里我们预测未来 N 个时间步后的收益率方向或价格变化。
# 预测未来 1 小时的收盘价变化
future_horizon = 60 # 假设 1 分钟数据，60 步即 1 小时
df_processed['target'] = df_processed.groupby('symbol')['close'].shift(-future_horizon)
df_processed = df_processed.dropna()
TFT 结合了可变选择网络、静态协变量处理和时序注意力机制。以下是核心组件的实现思路。
该部分用于确定哪些输入变量对特定时间步的预测最重要。通过门控机制抑制不相关的特征。
将输入变量映射到潜在空间，通常使用多层感知机 (MLP)。
利用自注意力机制捕捉长距离依赖关系，类似于标准 Transformer 中的 Encoder 层。
以下是一个简化的 PyTorch 实现框架，展示如何构建 TFT 的核心结构：
class TFTModel(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, num_layers=2):
super(TFTModel, self).__init__()
self.input_layer = nn.Linear(input_dim, hidden_dim)
self.attention = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=4, batch_first=True)
self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers=num_layers, batch_first=True)
self.fc = nn.Linear(hidden_dim, output_dim)
self.dropout = nn.Dropout(0.2)
def forward(self, x):
# x shape: [batch_size, seq_len, input_dim]
x = self.input_layer(x)
x = self.dropout(x)
# 时序注意力
attn_out, _ = self.attention(x, x, x)
x = x + attn_out # 残差连接
# LSTM 处理序列
lstm_out, _ = self.lstm(x)
x = lstm_out[:, -1, :] # 取最后一个时间步的输出
output = self.fc(x)
return output
将数据划分为训练集、验证集和测试集，注意保持时间顺序，避免数据泄露。
train_size = int(len(df_processed) * 0.7)
val_size = int(len(df_processed) * 0.15)
train_data = df_processed.iloc[:train_size]
val_data = df_processed.iloc[train_size:train_size+val_size]
test_data = df_processed.iloc[train_size+val_size:]
使用均方误差 (MSE) 作为回归任务的损失函数，配合 Adam 优化器。
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(50):
model.train()
optimizer.zero_grad()
outputs = model(train_inputs)
loss = criterion(outputs, train_targets)
loss.backward()
optimizer.step()
if (epoch + 1) % 10 == 0:
print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")
除了 MSE，我们还关注平均绝对误差 (MAE) 和方向准确率 (Directional Accuracy)，后者衡量预测涨跌方向的准确性。
from sklearn.metrics import mean_absolute_error, mean_squared_error
with torch.no_grad():
val_outputs = model(val_inputs)
mae = mean_absolute_error(val_targets.numpy(), val_outputs.numpy())
rmse = np.sqrt(mean_squared_error(val_targets.numpy(), val_outputs.numpy()))
print(f"Validation MAE: {mae:.4f}, RMSE: {rmse:.4f}")
本文详细阐述了利用时间融合 Transformer 模型进行股票价格预测的完整流程。从高频数据的清洗与特征工程，到 TFT 模型的架构设计与代码实现，再到训练评估策略，提供了一个可复现的技术方案。
TFT 模型的优势在于其能够同时处理静态协变量（如公司基本面）和动态协变量（如实时行情），并通过注意力机制提供一定的可解释性。在实际应用中，建议结合更多外部数据源（如宏观经济指标、新闻情感分析）进一步优化模型表现。此外，务必注意过拟合风险和回测中的未来函数陷阱，确保策略在实盘环境中的稳健性。
未来的工作可以集中在集成学习、强化学习策略优化以及更细粒度的风险建模上。希望本教程能为量化分析师和机器学习爱好者提供有价值的参考。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online