Transformer 时间融合模型在股票价格预测中的应用与代码实现

输出显示数据框包含超过 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)}")

探索性数据分析 (EDA) 与特征工程

在进行模型训练之前，深入理解数据分布和潜在模式至关重要。我们将进行以下特征工程步骤：

1. 基础统计特征

计算滚动窗口统计量，如移动平均线和波动率，这些是技术分析中的常用指标。

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)

2. 滞后特征 (Lag Features)

引入过去的时间步作为输入特征，帮助模型捕捉时间依赖性。

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)

3. 标签构建

定义预测目标。这里我们预测未来 N 个时间步后的收益率方向或价格变化。

# 预测未来 1 小时的收盘价变化
future_horizon = 60  # 假设 1 分钟数据，60 步即 1 小时
df_processed['target'] = df_processed.groupby('symbol')['close'].shift(-future_horizon)
df_processed = df_processed.dropna()

时间融合 Transformer (TFT) 模型架构

TFT 结合了可变选择网络、静态协变量处理和时序注意力机制。以下是核心组件的实现思路。

1. 变量选择网络

该部分用于确定哪些输入变量对特定时间步的预测最重要。通过门控机制抑制不相关的特征。

2. 观测值编码器

将输入变量映射到潜在空间，通常使用多层感知机 (MLP)。

3. 时序注意力机制

利用自注意力机制捕捉长距离依赖关系，类似于标准 Transformer 中的 Encoder 层。

4. 模型实现示例

以下是一个简化的 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

模型训练与评估

1. 数据集划分

将数据划分为训练集、验证集和测试集，注意保持时间顺序，避免数据泄露。

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:]

2. 损失函数与优化器

使用均方误差 (MSE) 作为回归任务的损失函数，配合 Adam 优化器。

criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

3. 训练循环

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}")

4. 评估指标

除了 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 模型的优势在于其能够同时处理静态协变量（如公司基本面）和动态协变量（如实时行情），并通过注意力机制提供一定的可解释性。在实际应用中，建议结合更多外部数据源（如宏观经济指标、新闻情感分析）进一步优化模型表现。此外，务必注意过拟合风险和回测中的未来函数陷阱，确保策略在实盘环境中的稳健性。

未来的工作可以集中在集成学习、强化学习策略优化以及更细粒度的风险建模上。希望本教程能为量化分析师和机器学习爱好者提供有价值的参考。